Updates for outputing WIC results

scoutbot/__init__.py

@@ -111,7 +111,7 @@ def pipeline(
         filepath (str): image filepath (relative or absolute)
 
     Returns:
-        list ( dict ): list of predictions
+        tuple ( float, list ( dict ) ): wic score, list of predictions
     """
     import utool as ut
 
@@ -122,6 +122,7 @@ def pipeline(
     wic_outputs = wic.post(wic.predict(wic.pre(tile_filepaths)))
 
     # Threshold for WIC
+    wic_ = max(wic_output.get('positive') for wic_output in wic_outputs)
     flags = [wic_output.get('positive') >= wic_thresh for wic_output in wic_outputs]
     loc_tile_grids = ut.compress(tile_grids, flags)
     loc_tile_filepaths = ut.compress(tile_filepaths, flags)
@@ -150,7 +151,7 @@ def pipeline(
             if exists(tile_filepath):
                 ut.delete(tile_filepath, verbose=False)
 
-    return detects
+    return wic_, detects
 
 
 def batch(
@@ -185,7 +186,7 @@ def batch(
         filepaths (list): list of str image filepath (relative or absolute)
 
     Returns:
-        list ( list ( dict ) ) : corresponding list of lists of predictions
+        tuple ( list ( float ), list ( list ( dict ) ) : corresponding list of wic scores, corresponding list of lists of predictions
     """
     import utool as ut
 
@@ -219,6 +220,14 @@ def batch(
 
     wic_outputs = wic.post(wic.predict(wic.pre(tile_filepaths)))
 
+    wic_dict = {}
+    for tile_img_filepath, wic_output in zip(tile_img_filepaths, wic_outputs):
+        wic_ = wic_output.get('positive')
+        existing_wic_ = wic_dict.get(tile_img_filepath, None)
+        if existing_wic_ is None:
+            existing_wic_ = wic_
+        wic_dict[tile_img_filepath] = max(existing_wic_, wic_)
+
     # 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)
@@ -242,9 +251,11 @@ def batch(
         batch[filepath]['loc']['outputs'].append(loc_output)
 
     # Run Aggregation
+    wic_list = []
     detects_list = []
     for filepath in filepaths:
         data = batch[filepath]
+        wic_ = wic_dict.get(filepath, None)
 
         img_shape = data['shape']
         loc_tile_grids = data['loc']['grids']
@@ -258,6 +269,8 @@ def batch(
             agg_thresh=agg_thresh,
             nms_thresh=agg_nms_thresh,
         )
+
+        wic_list.append(wic_)
         detects_list.append(detects)
 
     if clean:
@@ -265,7 +278,7 @@ def batch(
             if exists(tile_filepath):
                 ut.delete(tile_filepath, verbose=False)
 
-    return detects_list
+    return wic_list, detects_list
 
 
 def example():
@@ -286,6 +299,6 @@ def example():
 
     log.info(f'Running pipeline on image: {img_filepath}')
 
-    detects = pipeline(img_filepath)
+    wic_, detects = pipeline(img_filepath)
 
     log.info(ut.repr3(detects))

scoutbot/loc/__init__.py

@@ -31,26 +31,50 @@ from scoutbot.loc.transforms import (
 
 PWD = Path(__file__).absolute().parent
 
-BATCH_SIZE = 16
-INPUT_SIZE = (416, 416)
-INPUT_SIZE_H, INPUT_SIZE_W = INPUT_SIZE
-NETWORK_SIZE = (INPUT_SIZE_H, INPUT_SIZE_W, 3)
-
-NUM_CLASSES = 1
-ANCHORS = [
-    (1.3221, 1.73145),
-    (3.19275, 4.00944),
-    (5.05587, 8.09892),
-    (9.47112, 4.84053),
-    (11.2364, 10.0071),
-]
-CLASS_LABEL_MAP = ['elephant_savanna']
-LOC_THRESH = 0.4
-NMS_THRESH = 0.8
-
-ONNX_MODEL = 'scout.loc.5fbfff26.0.onnx'
-ONNX_MODEL_PATH = join(PWD, 'models', 'onnx', ONNX_MODEL)
-ONNX_MODEL_HASH = '85a9378311d42b5143f74570136f32f50bf97c548135921b178b46ba7612b216'
+PHASE1 = True
+
+if PHASE1:
+    BATCH_SIZE = 16
+    INPUT_SIZE = (416, 416)
+    INPUT_SIZE_H, INPUT_SIZE_W = INPUT_SIZE
+    NETWORK_SIZE = (INPUT_SIZE_H, INPUT_SIZE_W, 3)
+
+    NUM_CLASSES = 1
+    ANCHORS = [
+        (1.3221, 1.73145),
+        (3.19275, 4.00944),
+        (5.05587, 8.09892),
+        (9.47112, 4.84053),
+        (11.2364, 10.0071),
+    ]
+    CLASS_LABEL_MAP = ['elephant_savanna']
+    LOC_THRESH = 0.4
+    NMS_THRESH = 0.8
+
+    ONNX_MODEL = 'scout.loc.5fbfff26.0.onnx'
+    ONNX_MODEL_PATH = join(PWD, 'models', 'onnx', ONNX_MODEL)
+    ONNX_MODEL_HASH = '85a9378311d42b5143f74570136f32f50bf97c548135921b178b46ba7612b216'
+else:
+    BATCH_SIZE = 16
+    INPUT_SIZE = (416, 416)
+    INPUT_SIZE_H, INPUT_SIZE_W = INPUT_SIZE
+    NETWORK_SIZE = (INPUT_SIZE_H, INPUT_SIZE_W, 3)
+
+    NUM_CLASSES = 1
+    ANCHORS = [
+        (1.3221, 1.73145),
+        (3.19275, 4.00944),
+        (5.05587, 8.09892),
+        (9.47112, 4.84053),
+        (11.2364, 10.0071),
+    ]
+    CLASS_LABEL_MAP = ['elephant_savanna']
+    LOC_THRESH = 0.4
+    NMS_THRESH = 0.8
+
+    ONNX_MODEL = 'scout.loc.5fbfff26.0.onnx'
+    ONNX_MODEL_PATH = join(PWD, 'models', 'onnx', ONNX_MODEL)
+    ONNX_MODEL_HASH = '85a9378311d42b5143f74570136f32f50bf97c548135921b178b46ba7612b216'
 
 
 def fetch(pull=False):

scoutbot/scoutbot.py

@@ -83,7 +83,7 @@ def pipeline(
     agg_thresh /= 100.0
     agg_nms_thresh /= 100.0
 
-    detects = scoutbot.pipeline(
+    wic_, detects = scoutbot.pipeline(
         filepath,
         wic_thresh=wic_thresh,
         loc_thresh=loc_thresh,
@@ -94,9 +94,17 @@ def pipeline(
 
     if output:
         with open(output, 'w') as outfile:
-            json.dump(detects, outfile)
+            data = {
+                filepath: {
+                    'wic': wic_,
+                    'loc': detects,
+                }
+            }
+            json.dump(data, outfile)
     else:
-        log.info(ut.repr3(detects))
+        log.info(filepath)
+        log.info(f'WIC: {wic_:0.04f}')
+        log.info('LOC: {}'.format(ut.repr3(detects)))
 
 
 @click.command('batch')
@@ -155,7 +163,7 @@ def batch(
 
     log.info(f'Running batch on {len(filepaths)} files...')
 
-    detects_list = scoutbot.batch(
+    wic_list, detects_list = scoutbot.batch(
         filepaths,
         wic_thresh=wic_thresh,
         loc_thresh=loc_thresh,
@@ -163,16 +171,22 @@ def batch(
         agg_thresh=agg_thresh,
         agg_nms_thresh=agg_nms_thresh,
     )
-    results = zip(filepaths, detects_list)
+    results = zip(filepaths, wic_list, detects_list)
 
     if output:
-        detects = dict(results)
         with open(output, 'w') as outfile:
-            json.dump(detects, outfile)
+            data = {}
+            for filepath, wic_, detects in results:
+                data[filepath] = {
+                    'wic': wic,
+                    'loc': detects,
+                }
+                json.dump(data, outfile)
     else:
-        for filepath, detects in results:
+        for filepath, wic_, detects in results:
             log.info(filepath)
-            log.info(ut.repr3(detects))
+            log.info(f'WIC: {wic_:0.04f}')
+            log.info('LOC: {}'.format(ut.repr3(detects)))
 
 
 @click.command('example')

scoutbot/wic/__init__.py

@@ -26,12 +26,21 @@ from scoutbot.wic.dataloader import (  # NOQA
 
 PWD = Path(__file__).absolute().parent
 
-ONNX_MODEL = 'scout.wic.5fbfff26.3.0.onnx'
-ONNX_MODEL_PATH = join(PWD, 'models', 'onnx', ONNX_MODEL)
-ONNX_MODEL_HASH = 'cbc7f381fa58504e03b6510245b6b2742d63049429337465d95663a6468df4c1'
-ONNX_CLASSES = ['negative', 'positive']
-
-WIC_THRESH = 0.2
+PHASE1 = True
+
+
+if PHASE1:
+    ONNX_MODEL = 'scout.wic.5fbfff26.3.0.onnx'
+    ONNX_MODEL_PATH = join(PWD, 'models', 'onnx', ONNX_MODEL)
+    ONNX_MODEL_HASH = 'cbc7f381fa58504e03b6510245b6b2742d63049429337465d95663a6468df4c1'
+    ONNX_CLASSES = ['negative', 'positive']
+    WIC_THRESH = 0.2
+else:
+    ONNX_MODEL = 'scout.wic.5fbfff26.3.0.onnx'
+    ONNX_MODEL_PATH = join(PWD, 'models', 'onnx', ONNX_MODEL)
+    ONNX_MODEL_HASH = 'cbc7f381fa58504e03b6510245b6b2742d63049429337465d95663a6468df4c1'
+    ONNX_CLASSES = ['negative', 'positive']
+    WIC_THRESH = 0.2
 
 
 def fetch(pull=False):

tests/test_scoutbot.py

@@ -12,7 +12,7 @@ def test_fetch():
 def test_pipeline():
     img_filepath = abspath(join('examples', '1be4d40a-6fd0-42ce-da6c-294e45781f41.jpg'))
 
-    detects = scoutbot.pipeline(img_filepath)
+    wic_, detects = scoutbot.pipeline(img_filepath)
     assert len(detects) == 3
 
     targets = [