Add batch processing API and CLI command

scoutbot/__init__.py

@@ -146,6 +146,113 @@ def pipeline(
     return detects
 
 
+def batch(
+    filepaths,
+    wic_thresh=wic.WIC_THRESH,
+    loc_thresh=loc.LOC_THRESH,
+    loc_nms_thresh=loc.NMS_THRESH,
+    agg_thresh=agg.AGG_THRESH,
+    agg_nms_thresh=agg.NMS_THRESH,
+):
+    """
+    Run the ML pipeline on a given batch of image filepaths and return the detections
+    in a corresponding list.  The output is a list of outputs matching the output of
+    :func:`scoutbot.pipeline`, except the processing is done in batch and is much faster.
+
+    The final output is a list of lists of dictionaries, each representing a
+    single detection.  Each dictionary has a structure with the following keys:
+
+        ::
+
+            {
+                'l': class_label (str)
+                'c': confidence (float)
+                'x': x_top_left (float)
+                'y': y_top_left (float)
+                'w': width (float)
+                'h': height (float)
+            }
+
+    Args:
+        filepaths (list): list of str image filepath (relative or absolute)
+
+    Returns:
+        list ( list ( dict ) ) : corresponding list of lists of predictions
+    """
+    import utool as ut
+
+    # Run tiling
+    batch = {}
+    for filepath in filepaths:
+        img_shape, tile_grids, tile_filepaths = tile.compute(filepath)
+        data = {
+            'shape': img_shape,
+            'grids': tile_grids,
+            'filepaths': tile_filepaths,
+            'loc': {
+                'grids': [],
+                'outputs': [],
+            },
+        }
+        batch[filepath] = data
+
+    # Run WIC
+    tile_img_filepaths = []
+    tile_grids = []
+    tile_filepaths = []
+    for filepath in filepaths:
+        data = batch[filepath]
+        grids = data['grids']
+        filepaths = data['filepaths']
+        assert len(grids) == len(filepaths)
+        tile_img_filepaths += [filepath] * len(grids)
+        tile_grids += grids
+        tile_filepaths += filepaths
+
+    wic_outputs = wic.post(wic.predict(wic.pre(tile_filepaths)))
+
+    # Threshold for WIC
+    flags = [wic_output.get('positive') >= wic_thresh for wic_output in wic_outputs]
+    loc_tile_img_filepaths = ut.compress(tile_img_filepaths, flags)
+    loc_tile_grids = ut.compress(tile_grids, flags)
+    loc_tile_filepaths = ut.compress(tile_filepaths, flags)
+
+    # Run localizer
+    loc_data, loc_sizes = loc.pre(loc_tile_filepaths)
+    loc_preds = loc.predict(loc_data)
+    loc_outputs = loc.post(
+        loc_preds, loc_sizes, loc_thresh=loc_thresh, nms_thresh=loc_nms_thresh
+    )
+    assert len(loc_tile_grids) == len(loc_outputs)
+
+    for filepath, loc_tile_grid, loc_output in zip(
+        loc_tile_img_filepaths, loc_tile_grids, loc_outputs
+    ):
+        batch[filepath]['loc']['grids'].append(loc_tile_grid)
+        batch[filepath]['loc']['outputs'].append(loc_output)
+
+    # Run Aggregation
+    detects_list = []
+    for filepath in filepaths:
+        data = batch[filepath]
+
+        img_shape = data['shape']
+        loc_tile_grids = data['loc']['grids']
+        loc_outputs = data['loc']['outputs']
+        assert len(loc_tile_grids) == len(loc_outputs)
+
+        detects = agg.compute(
+            img_shape,
+            loc_tile_grids,
+            loc_outputs,
+            agg_thresh=agg_thresh,
+            nms_thresh=agg_nms_thresh,
+        )
+        detects_list.append(detects)
+
+    return detects_list
+
+
 def example():
     TEST_IMAGE = 'scout.example.jpg'
     TEST_IMAGE_HASH = (

scoutbot/scoutbot.py

@@ -20,11 +20,18 @@ def pipeline_filepath_validator(ctx, param, value):
     return value
 
 
-@click.command()
-@click.option(
-    '--filepath',
-    help='Path to image',
-    required=True,
+@click.command('fetch')
+def fetch():
+    """
+    Fetch the required machine learning ONNX models for the WIC and LOC
+    """
+    scoutbot.fetch()
+
+
+@click.command('pipeline')
+@click.argument(
+    'filepath',
+    nargs=1,
     type=str,
     callback=pipeline_filepath_validator,
 )
@@ -92,12 +99,80 @@ def pipeline(
         log.info(ut.repr3(detects))
 
 
-@click.command('fetch')
-def fetch():
+@click.command()
+@click.argument(
+    'filepaths',
+    nargs=-1,
+    type=str,
+)
+@click.option(
+    '--output',
+    help='Path to output JSON (if unspecified, results are printed to screen)',
+    default=None,
+    type=click.IntRange(0, 100, clamp=True),
+)
+@click.option(
+    '--wic_thresh',
+    help='Whole Image Classifier (WIC) confidence threshold',
+    default=wic.WIC_THRESH,
+    type=click.IntRange(0, 100, clamp=True),
+)
+@click.option(
+    '--loc_thresh',
+    help='Localizer (LOC) confidence threshold',
+    default=loc.LOC_THRESH,
+    type=click.IntRange(0, 100, clamp=True),
+)
+@click.option(
+    '--loc_nms_thresh',
+    help='Localizer (LOC) non-maximum suppression (NMS) threshold',
+    default=loc.NMS_THRESH,
+    type=click.IntRange(0, 100, clamp=True),
+)
+@click.option(
+    '--agg_thresh',
+    help='Aggregation (AGG) confidence threshold',
+    default=agg.AGG_THRESH,
+    type=click.IntRange(0, 100, clamp=True),
+)
+@click.option(
+    '--agg_nms_thresh',
+    help='Aggregation (AGG) non-maximum suppression (NMS) threshold',
+    default=agg.NMS_THRESH,
+    type=click.IntRange(0, 100, clamp=True),
+)
+def batch(
+    filepaths, output, wic_thresh, loc_thresh, loc_nms_thresh, agg_thresh, agg_nms_thresh
+):
     """
-    Fetch the required machine learning ONNX models for the WIC and LOC
+    Run the ScoutBot pipeline on an input image filepath
     """
-    scoutbot.fetch()
+    wic_thresh /= 100.0
+    loc_thresh /= 100.0
+    loc_nms_thresh /= 100.0
+    agg_thresh /= 100.0
+    agg_nms_thresh /= 100.0
+
+    log.info(f'Running batch on {len(filepaths)} files...')
+
+    detects_list = scoutbot.batch(
+        filepaths,
+        wic_thresh=wic_thresh,
+        loc_thresh=loc_thresh,
+        loc_nms_thresh=loc_nms_thresh,
+        agg_thresh=agg_thresh,
+        agg_nms_thresh=agg_nms_thresh,
+    )
+    results = zip(filepaths, detects_list)
+
+    if output:
+        detects = dict(results)
+        with open(output, 'w') as outfile:
+            json.dump(detects, outfile)
+    else:
+        for filepath, detects in results:
+            log.info(filepath)
+            log.info(ut.repr3(detects))
 
 
 @click.command('example')
@@ -118,6 +193,7 @@ def cli():
 
 cli.add_command(fetch)
 cli.add_command(pipeline)
+cli.add_command(batch)
 cli.add_command(example)