Merge pull request #19 from HakaiInstitute/release/0.2.0

Release 0.2.0
HakaiInstitute · Jun 14, 2022 · d8b3ddf · d8b3ddf
2 parents 20f4e6a + 7032191
commit d8b3ddf
Show file tree

Hide file tree

Showing 12 changed files with 112 additions and 110 deletions.
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/conda/meta.yaml b/conda/meta.yaml
@@ -22,12 +22,12 @@ requirements:
     - pip
   run:
     - python ~=3.7,<3.10
-    - fire ~=0.4
     - numpy ~=1.16
     - pytorch ~=1.10
     - rasterio ~=1.2
     - torchvision ~=0.11
     - tqdm ~=4.62
+    - typer ~=0.4
 
 test:
   imports:

diff --git a/docs/cli.rst b/docs/cli.rst
@@ -8,20 +8,17 @@ that the ``hakai-segmentation`` package is installed to.
 
     $ kom --help
 
-    NAME
-        kom
+    Usage: kom [OPTIONS] COMMAND [ARGS]...
 
-    SYNOPSIS
-        kom COMMAND
+    Options:
+      --install-completion  Install completion for the current shell.
+      --show-completion     Show completion for the current shell, to copy it or
+                            customize the installation.
+      --help                Show this message and exit.
 
-    COMMANDS
-        COMMAND is one of the following:
-
-         find-mussels
-           Detect mussels in source image and output the resulting classification raster to dest.
-
-         find-kelp
-           Detect kelp in source image and output the resulting classification raster to dest.
+    Commands:
+      find-kelp     Detect kelp in image at path `source` and output the...
+      find-mussels  Detect mussels in image at path `source` and output the...
 
 find-kelp
 ---------
@@ -30,43 +27,26 @@ find-kelp
 
     $ kom find-kelp --help
 
-    NAME
-        kom find-kelp - Detect kelp in source image and output the resulting classification raster to dest.
-
-    SYNOPSIS
-        kom find-kelp SOURCE DEST <flags>
-
-    DESCRIPTION
-        Detect kelp in source image and output the resulting classification raster to dest.
-
-    POSITIONAL ARGUMENTS
-        SOURCE
-            Type: str
-            Input image with Byte data type.
-        DEST
-            Type: str
-            File path location to save output to.
-
-    FLAGS
-        --crop_size=CROP_SIZE
-            Type: int
-            Default: 256
-            The size of cropped image square run through the segmentation model.
-        --padding=PADDING
-            Type: int
-            Default: 128
-            The number of context pixels added to each side of the cropped image squares.
-        --batch_size=BATCH_SIZE
-            Type: int
-            Default: 2
-            The batch size of cropped image sections to process together.
-        --no_gpu=NO_GPU
-            Type: bool
-            Default: False
-            Disable Cuda GPU usage and run on CPU only.
-
-    NOTES
-        You can also use flags syntax for POSITIONAL ARGUMENTS
+    Usage: kom find-kelp [OPTIONS] SOURCE DEST
+
+      Detect kelp in image at path `source` and output the resulting
+      classification raster to file at path `dest`.
+
+    Arguments:
+      SOURCE  Input image with Byte data type.  [required]
+      DEST    File path location to save output to.  [required]
+
+    Options:
+      --species / --presence  Segment to species or presence/absence level.
+                              [default: presence]
+      --crop-size INTEGER     The size for the cropped image squares run through
+                              the segmentation model.  [default: 256]
+      --padding INTEGER       The number of context pixels added to each side of
+                              the image crops.  [default: 128]
+      --batch-size INTEGER    The batch size of cropped image sections to process
+                              together.  [default: 2]
+      --gpu / --no-gpu        Enable or disable GPU, if available.  [default: gpu]
+      --help                  Show this message and exit.
 
 **Example usage**
 
@@ -82,44 +62,24 @@ find-mussels
 
     $ kom find-mussels --help
 
-    NAME
-        kom find-mussels - Detect mussels in source image and output the resulting classification raster to dest.
-
-    SYNOPSIS
-        kom find-kelp SOURCE DEST <flags>
-
-    DESCRIPTION
-        Detect kelp in source image and output the resulting classification raster to dest.
-
-    POSITIONAL ARGUMENTS
-        SOURCE
-            Type: str
-            Input image with Byte data type.
-        DEST
-            Type: str
-            File path location to save output to.
-
-    FLAGS
-        --crop_size=CROP_SIZE
-            Type: int
-            Default: 256
-            The size of cropped image square run through the segmentation model.
-        --padding=PADDING
-            Type: int
-            Default: 128
-            The number of context pixels added to each side of the cropped image squares.
-        --batch_size=BATCH_SIZE
-            Type: int
-            Default: 2
-            The batch size of cropped image sections to process together.
-        --no_gpu=NO_GPU
-            Type: bool
-            Default: False
-            Disable Cuda GPU usage and run on CPU only.
-
-    NOTES
-        You can also use flags syntax for POSITIONAL ARGUMENTS
+    Usage: kom find-mussels [OPTIONS] SOURCE DEST
+
+      Detect mussels in image at path `source` and output the resulting
+      classification raster to file at path `dest`.
+
+    Arguments:
+      SOURCE  Input image with Byte data type.  [required]
+      DEST    File path location to save output to.  [required]
 
+    Options:
+      --crop-size INTEGER   The size for the cropped image squares run through the
+                            segmentation model.  [default: 256]
+      --padding INTEGER     The number of context pixels added to each side of the
+                            image crops.  [default: 128]
+      --batch-size INTEGER  The batch size of cropped image sections to process
+                            together.  [default: 2]
+      --gpu / --no-gpu      Enable or disable GPU, if available.  [default: gpu]
+      --help                Show this message and exit.
 
 **Example usage**
 

diff --git a/environment.yml b/environment.yml
@@ -4,7 +4,7 @@ channels:
   - conda-forge
 dependencies:
   - python~=3.7,<3.10
-  - conda-forge::fire~=0.4
+  - conda-forge::typer~=0.4.1
   - conda-forge::numpy~=1.16
   - conda-forge::rasterio~=1.2
   - conda-forge::tqdm~=4.62

diff --git a/hakai_segmentation/cli.py b/hakai_segmentation/cli.py
@@ -1,14 +1,37 @@
-import fire
+import typer
 
-from hakai_segmentation.lib import find_kelp, find_mussels
+from hakai_segmentation import lib
 
+cli = typer.Typer()
 
-def cli():
-    """Run the python-fire CLI."""
-    fire.Fire({
-        "find-mussels": find_mussels,
-        "find-kelp": find_kelp
-    })
+
+@cli.command()
+def find_kelp(
+        source: str = typer.Argument(..., help="Input image with Byte data type."),
+        dest: str = typer.Argument(..., help="File path location to save output to."),
+        species: bool = typer.Option(False, "--species/--presence", help="Segment to species or presence/absence level."),
+        crop_size: int = typer.Option(256, help="The size for the cropped image squares run through the segmentation model."),
+        padding: int = typer.Option(128, help="The number of context pixels added to each side of the image crops."),
+        batch_size: int = typer.Option(2, help="The batch size of cropped image sections to process together."),
+        use_gpu: bool = typer.Option(True, "--gpu/--no-gpu", help="Enable or disable GPU, if available.")
+):
+    """Detect kelp in image at path `source` and output the resulting classification raster to file at path `dest`."""
+    lib.find_kelp(source=source, dest=dest, species=species,
+                  crop_size=crop_size, padding=padding, batch_size=batch_size, use_gpu=use_gpu)
+
+
+@cli.command()
+def find_mussels(
+        source: str = typer.Argument(..., help="Input image with Byte data type."),
+        dest: str = typer.Argument(..., help="File path location to save output to."),
+        crop_size: int = typer.Option(256, help="The size for the cropped image squares run through the segmentation model."),
+        padding: int = typer.Option(128, help="The number of context pixels added to each side of the image crops."),
+        batch_size: int = typer.Option(2, help="The batch size of cropped image sections to process together."),
+        use_gpu: bool = typer.Option(True, "--gpu/--no-gpu", help="Enable or disable GPU, if available.")
+):
+    """Detect mussels in image at path `source` and output the resulting classification raster to file at path `dest`."""
+    lib.find_mussels(source=source, dest=dest,
+                     crop_size=crop_size, padding=padding, batch_size=batch_size, use_gpu=use_gpu)
 
 
 if __name__ == '__main__':

diff --git a/hakai_segmentation/data/LRASPP_MobileNetV3_kelp_species_jit.pt b/hakai_segmentation/data/LRASPP_MobileNetV3_kelp_species_jit.pt
diff --git a/hakai_segmentation/data/__init__.py b/hakai_segmentation/data/__init__.py
@@ -4,3 +4,5 @@
     __name__, 'LRASPP_MobileNetV3_kelp_presence_jit.pt')
 lraspp_mussel_presence_torchscript_path = resource_filename(
     __name__, 'LRASPP_MobileNetV3_mussel_presence_jit.pt')
+lraspp_kelp_species_torchscript_path = resource_filename(
+    __name__, 'LRASPP_MobileNetV3_kelp_species_jit.pt')
diff --git a/hakai_segmentation/lib.py b/hakai_segmentation/lib.py
@@ -1,35 +1,39 @@
 from hakai_segmentation.managers import GeotiffSegmentation
-from hakai_segmentation.models import KelpPresenceSegmentationModel, MusselPresenceSegmentationModel
+from hakai_segmentation.models import KelpPresenceSegmentationModel, KelpSpeciesSegmentationModel, \
+    MusselPresenceSegmentationModel
 
 
-def find_kelp(source: str, dest: str,
-              crop_size: int = 256, padding: int = 128, batch_size: int = 2, no_gpu: bool = False):
+def find_kelp(source: str, dest: str, species: bool = False,
+              crop_size: int = 256, padding: int = 128, batch_size: int = 2, use_gpu: bool = True):
     """Detect kelp in image at path `source` and output the resulting classification raster to file at path `dest`.
 
     :param source: Input image with Byte data type.
     :param dest: File path location to save output to.
+    :param species: Do species classification instead of presence/absence.
     :param crop_size: The size of cropped image square run through the segmentation model.
     :param padding: The number of context pixels added to each side of the cropped image squares.
     :param batch_size: The batch size of cropped image sections to process together.
-    :param no_gpu: Disable Cuda GPU usage and run on CPU only.
+    :param use_gpu: Disable Cuda GPU usage and run on CPU only.
     """
-    model = KelpPresenceSegmentationModel(no_gpu=no_gpu)
+    if species:
+        model = KelpSpeciesSegmentationModel(use_gpu=use_gpu)
+    else:
+        model = KelpPresenceSegmentationModel(use_gpu=use_gpu)
     GeotiffSegmentation(model, source, dest,
                         crop_size=crop_size, padding=padding, batch_size=batch_size)()
 
 
 def find_mussels(source: str, dest: str,
-                 crop_size: int = 256, padding: int = 128, batch_size: int = 2, no_gpu: bool = False):
+                 crop_size: int = 256, padding: int = 128, batch_size: int = 2, use_gpu: bool = True):
     """Detect mussels in image at path `source` and output the resulting classification raster to file at path `dest`.
 
     :param source: Input image with Byte data type.
     :param dest: File path location to save output to.
     :param crop_size: The size of cropped image square run through the segmentation model.
     :param padding: The number of context pixels added to each side of the cropped image squares.
     :param batch_size: The batch size of cropped image sections to process together.
-    :param no_gpu: Disable Cuda GPU usage and run on CPU only.
+    :param use_gpu: Disable Cuda GPU usage and run on CPU only.
     """
-
-    model = MusselPresenceSegmentationModel(no_gpu=no_gpu)
+    model = MusselPresenceSegmentationModel(use_gpu=use_gpu)
     GeotiffSegmentation(model, source, dest,
                         crop_size=crop_size, padding=padding, batch_size=batch_size)()
diff --git a/hakai_segmentation/managers.py b/hakai_segmentation/managers.py
@@ -100,7 +100,6 @@ def on_start(self):
 
         # Setup progress bar
         self.progress = tqdm(
-            colour="CYAN",
             total=len(self.reader),
             desc="Processing"
         )

diff --git a/hakai_segmentation/models.py b/hakai_segmentation/models.py
@@ -2,12 +2,13 @@
 
 import torch
 
-from hakai_segmentation.data import lraspp_kelp_presence_torchscript_path, lraspp_mussel_presence_torchscript_path
+from hakai_segmentation.data import lraspp_kelp_presence_torchscript_path, lraspp_kelp_species_torchscript_path,\
+    lraspp_mussel_presence_torchscript_path
 
 
 class _Model(ABC):
-    def __init__(self, no_gpu: bool = False):
-        self.device = torch.device('cuda') if torch.cuda.is_available() and not no_gpu else torch.device('cpu')
+    def __init__(self, use_gpu: bool = True):
+        self.device = torch.device('cuda') if torch.cuda.is_available() and use_gpu else torch.device('cpu')
         self.model = self.load_model()
 
     @abstractmethod
@@ -34,6 +35,8 @@ def load_model(self) -> 'torch.nn.Module':
 class KelpPresenceSegmentationModel(_JITModel):
     torchscript_path = lraspp_kelp_presence_torchscript_path
 
+class KelpSpeciesSegmentationModel(_JITModel):
+    torchscript_path = lraspp_kelp_species_torchscript_path
 
 class MusselPresenceSegmentationModel(_JITModel):
     torchscript_path = lraspp_mussel_presence_torchscript_path
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 from setuptools import setup
 
 PACKAGE_NAME = 'hakai-segmentation'
-VERSION = '0.1.9'
+VERSION = '0.2.0'
 
 packages = [
     'hakai_segmentation',