diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100644 index 0000000..0b5a88a --- /dev/null +++ b/.idea/misc.xml @@ -0,0 +1,4 @@ + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..d9b6c13 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/conda/meta.yaml b/conda/meta.yaml index 3b06907..eff3edc 100644 --- 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 index 2db82f3..cdc39fc 100644 --- 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 - - 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 - - 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 index 4516d80..8883929 100644 --- 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 index 6bf1dc9..2ecd8ab 100644 --- 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 new file mode 100644 index 0000000..08fb2b2 Binary files /dev/null and b/hakai_segmentation/data/LRASPP_MobileNetV3_kelp_species_jit.pt differ diff --git a/hakai_segmentation/data/__init__.py b/hakai_segmentation/data/__init__.py index 1feb54a..e06a298 100644 --- 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') \ No newline at end of file diff --git a/hakai_segmentation/lib.py b/hakai_segmentation/lib.py index 57e7df5..b3a951c 100644 --- a/hakai_segmentation/lib.py +++ b/hakai_segmentation/lib.py @@ -1,25 +1,30 @@ 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. @@ -27,9 +32,8 @@ def find_mussels(source: str, dest: str, :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 index f1d0e8f..a44679a 100644 --- 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 index b1e7832..f0e108d 100644 --- 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 index bc26d2a..0bdb003 100644 --- 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',