Merge pull request #23 from WildMeOrg/add-classes-method

Add classes method

Author: TanyaStere42

scoutbot/__init__.py

@@ -62,6 +62,7 @@
 
 
 from scoutbot import agg, loc, tile, wic, tile_batched  # NOQA
+from scoutbot.loc import CONFIGS as LOC_CONFIGS # NOQA
 
 # from tile_batched.models import Yolov8DetectionModel
 # from tile_batched import get_sliced_prediction_batched
@@ -460,6 +461,26 @@ def batch_v3(
     return wic_list, detects_list
 
 
+def get_classes():
+    classes = set()
+    for config in ['v3', 'v3-cls']:
+        yolov8_model_path = loc.fetch(config=config)
+        model = tile_batched.Yolov8DetectionModel(
+            model_path=yolov8_model_path,
+            confidence_threshold=0.5,
+            device='cpu',
+        )
+        model_classes = list(model.category_names)
+        classes.update(model_classes)
+
+    for config in ['phase1', 'mvp']:
+        model_classes = LOC_CONFIGS[config]['classes']
+        classes.update(model_classes)
+
+    classes = list(classes)
+    return classes
+
+
 def example():
     """
     Run the pipeline on an example image with the default configuration

scoutbot/scoutbot.py

@@ -29,13 +29,16 @@ def pipeline_filepath_validator(ctx, param, value):
     ),
 ]
 
-shared_options = [
+output_option = [
     click.option(
         '--output',
         help='Path to output JSON (if unspecified, results are printed to screen)',
         default=None,
         type=str,
     ),
+]
+
+shared_options = [
     click.option(
         '--backend_device',  # torch backend device
         help='Specifies the device for inference (see YOLO and PyTorch documentation for more information).',
@@ -93,6 +96,7 @@ def _add_options(func):
 )
 @add_options(model_option)
 @add_options(shared_options)
+@add_options(output_option)
 def pipeline(
     filepath,
     config,
@@ -183,6 +187,7 @@ def pipeline(
 )
 @add_options(model_option)
 @add_options(shared_options)
+@add_options(output_option)
 def batch(
     filepaths,
     config,
@@ -300,6 +305,21 @@ def example():
     scoutbot.example()
 
 
+@click.command('get_classes')
+@add_options(output_option)
+def get_classes(output):
+    """
+    Run a test of the pipeline on an example image with the default configuration.
+    """
+    classes = scoutbot.get_classes()
+    log.debug('Outputting classes list...')
+    if output:
+        with open(output, 'w') as outfile:
+            json.dump(classes, outfile)
+    else:
+        print(ut.repr3(classes))
+
+
 @click.group()
 def cli():
     """
@@ -312,6 +332,8 @@ def cli():
 cli.add_command(batch)
 cli.add_command(fetch)
 cli.add_command(example)
+cli.add_command(get_classes)
+
 
 if __name__ == '__main__':
     cli()