Refactor scripts into a CLI

README.md

@@ -158,6 +158,199 @@ python3 wing-segmentation/landmark_scripts/create_wing_folders.py --input_dir /p
 ```
 
 **Flip images**
-```
+```console
 python3 wing-segmentation/landmark_scripts/flip_images_horizontally.py --input_dir /path/to/wing/category/folder
 ```
+
+# CLI Help
+
+> [!CAUTION]
+> The CLI is still under development and may not be fully functional. Please add an [issue](https://github.com/Imageomics/wing-segmentation/issues) for any bugs or feature requests.
+
+The wing segmentation CLI tool is designed for convenient and flexible segmentation of butterfly images.
+
+```console
+usage: wingseg [-h] {segment,scan-runs} ...
+
+Wing Segmenter CLI
+
+options:
+  -h, --help           show this help message and exit
+
+Commands:
+  {segment,scan-runs}
+    segment            Segment images and store segmentation masks.
+    scan-runs          List existing processing runs for a dataset.
+```
+## Options for `wingseg segment`
+
+This command segments images and stores segmentation masks with a variety of options for resizing, padding, background removal, and more.
+
+```console
+usage: wingseg segment [-h] --dataset DATASET [--size SIZE [SIZE ...]] [--resize-mode {distort,pad}] [--padding-color {black,white}]
+                       [--interpolation {nearest,linear,cubic,area,lanczos4,linear_exact,nearest_exact}] [--bbox-padding BBOX_PADDING]
+                       [--outputs-base-dir OUTPUTS_BASE_DIR | --custom-output-dir CUSTOM_OUTPUT_DIR] [--sam-model SAM_MODEL] [--yolo-model YOLO_MODEL]
+                       [--device {cpu,cuda}] [--visualize-segmentation] [--crop-by-class] [--force] [--remove-crops-background] [--remove-full-background]
+                       [--background-color {white,black}]
+
+options:
+  -h, --help            show this help message and exit
+  --dataset DATASET     Path to dataset images (default: None)
+  --outputs-base-dir OUTPUTS_BASE_DIR
+                        Base path to store outputs. (default: None)
+  --custom-output-dir CUSTOM_OUTPUT_DIR
+                        Fully custom directory to store all output files. (default: None)
+  --sam-model SAM_MODEL
+                        SAM model to use (e.g., facebook/sam-vit-base) (default: facebook/sam-vit-base)
+  --yolo-model YOLO_MODEL
+                        YOLO model to use (local path or Hugging Face repo). (default:
+                        imageomics/butterfly_segmentation_yolo_v8:yolov8m_shear_10.0_scale_0.5_translate_0.1_fliplr_0.0_best.pt)
+  --device {cpu,cuda}   Device to use for processing. (default: cpu)
+  --visualize-segmentation
+                        Generate and save segmentation visualizations. (default: False)
+  --crop-by-class       Enable cropping of segmented classes into crops/ directory. (default: False)
+  --force               Force reprocessing even if outputs already exist. (default: False)
+
+Resizing Options:
+  --size SIZE [SIZE ...]
+                        Target size. Provide one value for square dimensions or two for width and height. (default: None)
+  --resize-mode {distort,pad}
+                        Resizing mode. "distort" resizes without preserving aspect ratio, "pad" preserves aspect ratio and adds padding if necessary.
+                        (default: None)
+  --padding-color {black,white}
+                        Padding color to use when --resize-mode is "pad". (default: None)
+  --interpolation {nearest,linear,cubic,area,lanczos4,linear_exact,nearest_exact}
+                        Interpolation method to use when resizing. For upscaling, "lanczos4" is recommended. (default: area)
+
+Bounding Box Options:
+  --bbox-padding BBOX_PADDING
+                        Padding to add to bounding boxes in pixels. Defaults to no padding. (default: None)
+
+Background Removal Options:
+  --remove-crops-background
+                        Remove background from cropped images. (default: False)
+  --remove-full-background
+                        Remove background from the entire (resized or original) image. (default: False)
+  --background-color {white,black}
+                        Background color to use when removing background. (default: None)
+```
+
+## Options for `wingseg scan-runs`
+
+This command provides a tabular overview of segmentation runs for comparing effects of segmentation option settings:
+
+```console
+usage: wingseg scan-runs [-h] --dataset DATASET [--output-dir OUTPUT_DIR]
+
+options:
+  -h, --help            show this help message and exit
+  --dataset DATASET     Path to the dataset directory.
+  --output-dir OUTPUT_DIR
+                        Base path where outputs were stored.
+```
+
+Example usage:
+```console
+wingseg segment --dataset ../data/input/ \
+--outputs-base-dir ../data/output/ \
+--visualize-segmentation \
+--crop-by-class \
+--size 512 \
+--resize-mode pad \
+--padding-color white \
+--interpolation cubic \
+--remove-crops-background \
+--remove-full-background \
+--background-color white
+```
+Depending on the contents of `../data/input/`, the command above will produce the following status indicator:
+```console
+INFO:root:Loading YOLO model: imageomics/butterfly_segmentation_yolo_v8:yolov8m_shear_10.0_scale_0.5_translate_0.1_fliplr_0.0_best.pt
+INFO:root:YOLO model loaded onto cpu
+INFO:root:Loading SAM model: facebook/sam-vit-base
+INFO:root:Loaded SAM model and processor successfully.
+INFO:root:Processing 18 images
+INFO:root:Output directory: /abs/path/to/data/output/input_3354acb9-b295-5d07-9397-8ec5c74cee37
+Processing Images:   6%|█████▌                                                                                               | 1/18 [00:14<04:09, 14.67s/image]
+```
+Note that the unique identifier appended to the output directory is a UUID that depends on certain options specified in the command as well as the input dataset. This is to ensure that the output directory is unique and does not overwrite existing results.
+
+For example, it may be useful to compare the effects of resize dimensions with squares of size [256, 512, 1024].
+
+Once these are processed, you can use the `scan-runs` command for a tabular overview of the segmentation runs:
+```console
+wingseg scan-runs --dataset ../data/input/ --output-dir ../data/output/
+Found 3 processing runs for dataset 'input':
+
+                                             Processing Runs                                             
+┏━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━┓
+┃       ┃ Run UUID ┃           ┃            ┃             ┃ Resize  ┃               ┃          ┃        ┃
+┃ Run # ┃ Prefix   ┃ Completed ┃ Num Images ┃ Resize Dims ┃  Mode   ┃    Interp     ┃ BBox Pad ┃ Errors ┃
+┡━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━┩
+│     1 │ 3354acb9 │    Yes    │         18 │   512x512   │   pad   │     cubic     │        0 │  None  │
+├───────┼──────────┼───────────┼────────────┼─────────────┼─────────┼───────────────┼──────────┼────────┤
+│     2 │ 0f27d745 │    Yes    │         18 │   256x256   │   pad   │     cubic     │        0 │  None  │
+├───────┼──────────┼───────────┼────────────┼─────────────┼─────────┼───────────────┼──────────┼────────┤
+│     3 │ 8e9ae0a2 │    Yes    │         18 │  1024x1024  │   pad   │     cubic     │        0 │  None  │
+└───────┴──────────┴───────────┴────────────┴─────────────┴─────────┴───────────────┴──────────┴────────┘
+```
+
+This can be helpful navigating the outputs of multiple segmentation runs:
+```console
+$ ls -1 ../data/output/*
+../data/output/input_0f27d745-12ce-50b9-a28c-5641dbfaea49:
+crops
+crops_bkgd_removed
+full_bkgd_removed
+logs
+masks
+metadata
+resized
+seg_viz
+
+../data/output/input_3354acb9-b295-5d07-9397-8ec5c74cee37:
+<similarly>
+
+../data/output/input_8e9ae0a2-992c-579d-bb51-b8715442bcf4:
+<similarly>
+```
+
+Inspecting data in the `seg_viz/`, we can see that the 1024x1024 products have segmentation masks that differ from the 512x512 and 256x256 products.
+
+1024x1024 (Run 8e9ae0a2):
+
+![1024 resize segmentation result visualization](readme_images/STRI_WOM_0011_V_viz_1024.png)
+
+512x512 (Run 3354acb9):
+
+![512 resize segmentation result visualization](readme_images/STRI_WOM_0011_V_viz_512.png)
+
+256x256 (Run 0f27d745):
+
+![256 resize segmentation result visualization](readme_images/STRI_WOM_0011_V_viz_256.png)
+
+A potential fix for this could be to add padding to the bounding boxes (with the `--bbox-padding` option) wherever results are inconsistent with expectations.
+
+
+## Further Development
+To use and continue building the CLI features, set up and activate a virtual environment, and build interactively with `pip install -e .[dev]`.
+
+> [!TIP] 
+> [`uv`](https://github.com/astral-sh/uv) is a fast, Rust-based package manager.
+> If using on an HPC system, you may install `uv` into your user path or a conda environment (latter illustrated here).
+
+```console
+conda create -n uv -c conda-forge --solver=libmamba python=3.10 uv -y
+```
+```console
+conda activate uv
+```
+```console
+uv venv
+```
+```console
+source .venv/bin/activate # or source .venv/Scripts/activate on Windows
+```
+```console
+uv pip install -e .[dev]
+```

pyproject.toml

@@ -0,0 +1,57 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/wing_segmenter"]
+
+[project]
+name = "wing-segmenter"
+version = "0.1.0"
+description = "A CLI tool for Lepidopteran wing preprocessing and segmentation."
+authors = [
+    { name = "Michelle Ramirez", email = "[email protected]" },
+    { name = "Matthew J. Thompson", email = "[email protected]"}
+]
+license = { text = "MIT" }
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch",
+    "torchvision",
+    "numpy",
+    "pandas",
+    "opencv-python",
+    "Pillow",
+    "matplotlib",
+    "scikit-image",
+    "scikit-learn",
+    "ultralytics",
+    "rich",
+    "tqdm",
+    "huggingface-hub",
+    "pycocotools",
+    "wget",
+    "segment-anything",
+    "transformers",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "ruff",
+]
+
+[project.urls]
+Documentation = "https://github.com/Imageomics/wing-segmentation#readme"
+Issues = "https://github.com/Imageomics/wing-segmentation/issues"
+Source = "https://github.com/Imageomics/wing-segmentation/"
+
+[project.scripts]
+wingseg = "wing_segmenter.__main__:main"
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.version]
+path = "src/wing_segmenter/__init__.py"

src/wing_segmenter/__init__.py

@@ -0,0 +1,7 @@
+__version__ = "0.1.0"
+
+def __getattr__(name):
+    if name == 'Segmenter':
+        from wing_segmenter.segmenter import Segmenter
+        return Segmenter
+    raise AttributeError(f"module {__name__} has no attribute {name}")

src/wing_segmenter/__main__.py

@@ -0,0 +1,4 @@
+from wing_segmenter.cli import main
+
+if __name__ == "__main__":
+    main()

src/wing_segmenter/cli.py

@@ -0,0 +1,125 @@
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog='wingseg',
+        description="Wing Segmenter CLI",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    subparsers = parser.add_subparsers(title='Commands', dest='command', required=True)
+
+    # Subcommand: segment
+    segment_parser = subparsers.add_parser(
+        'segment',
+        help='Segment images and store segmentation masks.',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    # Required argument
+    segment_parser.add_argument('--dataset', required=True, help='Path to dataset images')
+
+    # Resizing options
+    resize_group = segment_parser.add_argument_group('Resizing Options')
+
+    # Dimension specifications
+    resize_group.add_argument('--size', nargs='+', type=int,
+                              help='Target size. Provide one value for square dimensions or two for width and height.')
+
+    # Resizing mode
+    resize_group.add_argument('--resize-mode', choices=['distort', 'pad'], default=None,
+                              help='Resizing mode. "distort" resizes without preserving aspect ratio, "pad" preserves aspect ratio and adds padding if necessary.')
+
+    # Padding options (to preserve aspect ratio)
+    resize_group.add_argument('--padding-color', choices=['black', 'white'], default=None,
+                              help='Padding color to use when --resize-mode is "pad".')
+
+    # Interpolation options
+    resize_group.add_argument('--interpolation', choices=['nearest', 'linear', 'cubic', 'area', 'lanczos4', 'linear_exact', 'nearest_exact'],
+                              default='area',
+                              help='Interpolation method to use when resizing. For upscaling, "lanczos4" is recommended.')
+
+    # Bounding box padding option
+    bbox_group = segment_parser.add_argument_group('Bounding Box Options')
+    bbox_group.add_argument('--bbox-padding', type=int, default=None,
+                            help='Padding to add to bounding boxes in pixels. Defaults to no padding.')
+
+
+    # Output options within mutually exclusive group
+    output_group = segment_parser.add_mutually_exclusive_group()
+    output_group.add_argument('--outputs-base-dir', default=None, help='Base path to store outputs.')
+    output_group.add_argument('--custom-output-dir', default=None, help='Fully custom directory to store all output files.')
+
+    # General processing options
+    segment_parser.add_argument('--sam-model', default='facebook/sam-vit-base',
+                                help='SAM model to use (e.g., facebook/sam-vit-base)')
+    segment_parser.add_argument('--yolo-model', default='imageomics/butterfly_segmentation_yolo_v8:yolov8m_shear_10.0_scale_0.5_translate_0.1_fliplr_0.0_best.pt',
+                                help='YOLO model to use (local path or Hugging Face repo).')
+    segment_parser.add_argument('--device', choices=['cpu', 'cuda'], default='cpu',
+                                help='Device to use for processing.')
+    segment_parser.add_argument('--visualize-segmentation', action='store_true',
+                                help='Generate and save segmentation visualizations.')
+    segment_parser.add_argument('--crop-by-class', action='store_true',
+                                help='Enable cropping of segmented classes into crops/ directory.')
+    segment_parser.add_argument('--force', action='store_true',
+                                help='Force reprocessing even if outputs already exist.')
+
+    # Background removal options
+    bg_group = segment_parser.add_argument_group('Background Removal Options')
+    bg_group.add_argument('--remove-crops-background', action='store_true',
+                           help='Remove background from cropped images.')
+    bg_group.add_argument('--remove-full-background', action='store_true',
+                           help='Remove background from the entire (resized or original) image.')
+    bg_group.add_argument('--background-color', choices=['white', 'black'], default=None,
+                           help='Background color to use when removing background.')
+
+    # Subcommand: scan-runs
+    scan_parser = subparsers.add_parser('scan-runs', help='List existing processing runs for a dataset.')
+    scan_parser.add_argument('--dataset', required=True, help='Path to the dataset directory.')
+    scan_parser.add_argument('--output-dir', default=None, help='Base path where outputs were stored.')
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Command input validations
+    if args.command == 'segment':
+        # If size is provided, enforce resizing options
+        if args.size:
+            if len(args.size) not in [1, 2]:
+                parser.error('--size must accept either one value (square resize) or two values (width and height).')
+            if not args.resize_mode:
+                parser.error('--resize-mode must be specified when --size is provided.')
+        # If no size is provided, ensure that resizing options were not explicitly set
+        else:
+            if args.resize_mode is not None:
+                parser.error('Resizing options (--resize-mode) require --size to be specified.')
+            if args.padding_color is not None:
+                parser.error('Resizing options (--padding-color) require --size to be specified.')
+
+        # --remove-crops-background requires --crop-by-class
+        if args.remove_crops_background and not args.crop_by_class:
+            parser.error('--remove-crops-background requires --crop-by-class to be set.')
+
+        # Need to set croped or full background removal to set background color
+        if args.background_color and not (args.remove_crops_background or args.remove_full_background):
+            parser.error('--background-color can only be set when background removal is enabled.')
+
+        # Ensure that if --custom-output-dir is set, --outputs-base-dir is not used
+        if args.custom_output_dir and args.outputs_base_dir:
+            parser.error('Cannot specify both --outputs-base-dir and --custom-output-dir. Choose one.')
+
+        # Validate bbox-padding
+        if args.bbox_padding is not None and args.bbox_padding < 0:
+            parser.error('--bbox-padding must be a non-negative integer.')
+
+    # Execute the subcommand
+    if args.command == 'segment':
+        from wing_segmenter.segmenter import Segmenter
+
+        segmenter = Segmenter(args)
+        segmenter.process_dataset()
+
+    elif args.command == 'scan-runs':
+        from wing_segmenter.run_scanner import scan_runs
+
+        scan_runs(dataset_path=args.dataset, output_base_dir=args.output_dir)