From f344865812ede812a4bff7b15fd24f71b0d8c418 Mon Sep 17 00:00:00 2001 From: danellecline Date: Mon, 29 Apr 2024 18:07:07 -0700 Subject: [PATCH] feat: added support for specifying start/ending file names to support skipping image blocks --- sdcat/cluster/commands.py | 20 +++++++++++++++++++- sdcat/common_args.py | 10 +++++++++- sdcat/detect/commands.py | 22 +++++++++++++++++++++- 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/sdcat/cluster/commands.py b/sdcat/cluster/commands.py index efe904e..51cc60e 100644 --- a/sdcat/cluster/commands.py +++ b/sdcat/cluster/commands.py @@ -22,6 +22,8 @@ @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) @@ -29,7 +31,7 @@ @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: diff --git a/sdcat/common_args.py b/sdcat/common_args.py index 9beee05..05b3f05 100644 --- a/sdcat/common_args.py +++ b/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') \ No newline at end of file + 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') \ No newline at end of file diff --git a/sdcat/detect/commands.py b/sdcat/detect/commands.py index eb16850..58f2446 100644 --- a/sdcat/detect/commands.py +++ b/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}')