feat: added support for specifying start/ending file names to support…

… skipping image blocks

sdcat/cluster/commands.py

@@ -22,14 +22,16 @@
 
 @click.command('cluster', help='Cluster detections. See cluster --config-ini to override cluster defaults.')
 @common_args.config_ini
+@common_args.start_image
+@common_args.end_image
 @click.option('--det-dir', help='Input folder(s) with raw detection results', multiple=True)
 @click.option('--save-dir', help='Output directory to save clustered detection results')
 @click.option('--device', help='Device to use.', type=int)
 @click.option('--alpha', help='Alpha is a parameter that controls the linkage.', type=float, default=0.92)
 @click.option('--cluster_selection_epsilon', help='Epsilon is a parameter that controls the linkage', type=float, default=0.0)
 @click.option('--min_cluster_size', help='The minimum number of samples in a group for that group to be considered a cluster', type=int, default=2)
 
-def run_cluster(det_dir, save_dir, device, config_ini, alpha, cluster_selection_epsilon, min_cluster_size):
+def run_cluster(det_dir, save_dir, device, config_ini, alpha, cluster_selection_epsilon, min_cluster_size, start_image, end_image):
     config = cfg.Config(config_ini)
     max_area = int(config('cluster', 'max_area'))
     min_area = int(config('cluster', 'min_area'))
@@ -87,6 +89,22 @@ def run_cluster(det_dir, save_dir, device, config_ini, alpha, cluster_selection_
         err(f'Found {df["image_path"].isnull().sum()} detections with no image_path')
         return
 
+    # If start_image is set, find the index of the start_image in the list of images
+    if start_image:
+        start_image = Path(start_image)
+        start_image = start_image.resolve()
+        start_image = start_image.stem
+        start_image_index = df[df['image_path'].str.contains(start_image)].index[0]
+        df = df.iloc[start_image_index:]
+
+    # If end_image is set, find the index of the end_image in the list of images
+    if end_image:
+        end_image = Path(end_image)
+        end_image = end_image.resolve()
+        end_image = end_image.stem
+        end_image_index = df[df['image_path'].str.contains(end_image)].index[0]
+        df = df.iloc[:end_image_index]
+
     # Filter by saliency, area, score or day/night
     size_before = len(df)
     if 'saliency' in df.columns:

sdcat/common_args.py

@@ -10,4 +10,12 @@
 config_ini = click.option('--config-ini',
                           type=str,
                           default=default_config_ini,
-                          help=f'Path to config file to override. Defaults are in {default_config_ini}. Copy to your own custom.ini file to override')
+                          help=f'Path to config file to override. Defaults are in {default_config_ini}. Copy to your own custom.ini file to override')
+
+start_image = click.option('--start-image',
+                            type=str,
+                            help='Start image name')
+
+end_image = click.option('--end-image',
+                            type=str,
+                            help='End image name')

sdcat/detect/commands.py

@@ -27,6 +27,8 @@
                     f'Use --skip-sahi or --skip-saliency to exclude. '
                     f'See --config-ini to override detection defaults in {default_config_ini}.')
 @common_args.config_ini
+@common_args.start_image
+@common_args.end_image
 @click.option('--show', is_flag=True, help='Show algorithm steps.')
 @click.option('--image-dir', required=True, help='Directory with images to run sliced detection.')
 @click.option('--save-dir', required=True, help='Save detections to this directory.')
@@ -45,7 +47,7 @@
 def run_detect(show: bool, image_dir: str, save_dir: str, model: str,
                slice_size_width: int, slice_size_height: int, scale_percent: int,
                device: str, conf: float, skip_sahi: bool, skip_saliency: bool, spec_remove: bool,
-               config_ini: str, clahe: bool):
+               config_ini: str, clahe: bool, start_image: str, end_image: str):
     config = cfg.Config(config_ini)
     max_area = int(config('detect', 'max_area'))
     min_area = int(config('detect', 'min_area'))
@@ -153,6 +155,24 @@ def run_detect(show: bool, image_dir: str, save_dir: str, model: str,
     images = [file for file in images_path.rglob('*')
               if file.as_posix().endswith(('jpeg', 'png', 'jpg', 'JPEG', 'PNG', 'JPG', 'tif', 'tiff'))]
 
+    # If start_image is set, find the index of the start_image in the list of images
+    if start_image:
+        start_image = Path(start_image)
+        start_image_index = next((i for i, image in enumerate(images) if start_image.name in image.name), None)
+        if start_image_index is None:
+            warn(f'Start image {start_image} not found in images')
+            return
+        images = images[start_image_index:]
+
+    # If end_image is set, find the index of the end_image in the list of images
+    if end_image:
+        end_image = Path(end_image)
+        end_image_index = next((i for i, image in enumerate(images) if end_image.name in image.name), None)
+        if end_image_index is None:
+            warn(f'End image {end_image} not found in images')
+            return
+        images = images[:end_image_index + 1]
+
     num_images = len(images)
     info(f'Found {num_images} images in {images_path}')