Refactor scripts into a CLI

README.md

@@ -161,3 +161,52 @@ python3 wing-segmentation/landmark_scripts/create_wing_folders.py --input_dir /p
 ```
 python3 wing-segmentation/landmark_scripts/flip_images_horizontally.py --input_dir /path/to/wing/category/folder
 ```
+
+# CLI Development
+
+```bash
+usage: ws [-h] {resize} ...
+
+Wing Segmenter CLI
+
+options:
+  -h, --help  show this help message and exit
+
+Commands:
+  {resize}
+    resize    Resize images and store them in a structured output directory.
+```
+## Optional preprocessing
+```bash
+usage: ws resize [-h] --source SOURCE [--output OUTPUT] --resize-dim WIDTH HEIGHT [--num-workers NUM_WORKERS]
+                 [--interpolation {nearest,linear,cubic,area,lanczos4,linear_exact,nearest_exact}]
+
+options:
+  -h, --help            show this help message and exit
+  --source SOURCE       Path to source images
+  --output OUTPUT       Base path to output resized images.
+                        Final path will include <WIDTH>x<HEIGHT>/<INTERPOLATION>.
+                        Default: <SOURCE>_resize/<WIDTH>x<HEIGHT>/<INTERPOLATION>, neighboring SOURCE.
+                        If SOURCE has nested directories, the output will mirror the structure.
+  --resize-dim WIDTH HEIGHT
+                        Resize dimensions (WIDTH HEIGHT)
+  --num-workers NUM_WORKERS
+                        Number of parallel workers (default: 1)
+  --interpolation {nearest,linear,cubic,area,lanczos4,linear_exact,nearest_exact}
+                        OpenCV interpolation method to use (default: area)
+```
+
+To use and continue building the CLI features:
+
+## Set up an environment and install the package.
+```bash
+conda create -n ws -c conda-forge --solver=libmamba python=3.10 uv -y
+```
+```bash
+conda activate ws
+```
+```bash
+uv pip install -e .[dev]
+```
+> Note: [`uv`](https://github.com/astral-sh/uv) is a fast, Rust-based package manager.
+

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,3 @@
+__version__ = "0.1.0"
+
+from wing_segmenter.segmenter import Segmenter

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,107 @@
+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
+    resize_group.add_argument('--padding-color', choices=['black', 'white'], default='black',
+                              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.')
+
+    # General processing options
+    segment_parser.add_argument('--output-dir', default=None, help='Base path to store outputs.')
+    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('--num-workers', type=int, default=1,
+                                help='Number of worker threads to use for processing.')
+    segment_parser.add_argument('--device', choices=['cpu', 'cuda'], default='cpu',
+                                help='Device to use for processing.')
+    segment_parser.add_argument('--save-intermediates', action='store_true',
+                                help='Save intermediate files (resized images and segmentation masks).')
+    segment_parser.add_argument('--visualize-segmentation', action='store_true',
+                                help='Generate and save segmentation visualizations.')
+    segment_parser.add_argument('--force', action='store_true',
+                                help='Force reprocessing even if outputs already exist.')
+    segment_parser.add_argument('--crop-by-class', action='store_true',
+                                help='Enable cropping of segmented classes into crops/ directory.')
+
+    # CLI Flags for background removal
+    bg_group = segment_parser.add_argument_group('Background Removal Options')
+    bg_group.add_argument('--remove-background', action='store_true', default=False,
+                           help='Remove background from cropped images.')
+    bg_group.add_argument('--background-color', choices=['white', 'black'], default='black',
+                           help='Background color to use when removing background. Applicable only if --remove-background is set.')
+    bg_group.add_argument('--remove-bg-full', action='store_true',
+                           help='Remove background from the entire (resized or original) image.')
+
+    # 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()
+
+    # Validation for resizing options
+    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 args.resize_mode == 'pad' and args.padding_color is None:
+                args.padding_color = 'black'
+        # If no size is provided, ensure that resizing options were not explicitly set
+        else:
+            if args.resize_mode is not None or args.padding_color != 'black':
+                parser.error('Resizing options (resize-mode, padding-color) require --size to be specified.')
+
+        # Additional validation for background removal flags
+        if (args.remove_background or args.remove_bg_full) and not args.crop_by_class:
+            parser.error('--remove-background and --remove-bg-full require --crop-by-class to be set.')
+
+    # 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)

src/wing_segmenter/constants.py

@@ -0,0 +1,12 @@
+CLASSES = {
+    0: 'background',
+    1: 'right_forewing',
+    2: 'left_forewing',
+    3: 'right_hindwing',
+    4: 'left_hindwing',
+    5: 'ruler',
+    6: 'white_balance',
+    7: 'label',
+    8: 'color_card',
+    9: 'body'
+}

src/wing_segmenter/exceptions.py

@@ -0,0 +1,11 @@
+class ModelLoadError(Exception):
+    """Exception raised when a model fails to load."""
+    pass
+
+class ImageProcessingError(Exception):
+    """Exception raised when an error occurs during image processing."""
+    pass
+
+class MetadataError(Exception):
+    """Exception raised when there is an error with metadata handling."""
+    pass