Config file

InferenceConfig.py

@@ -0,0 +1,73 @@
+from enum import Enum
+
+### Configuration options
+WEIGHTS = 'models/v5m_896_300best.pt'
+# will need to configure these based on GPU hardware
+BATCH_SIZE = 32
+
+CONF_THRES = 0.05 # detection
+NMS_IOU  = 0.2 # NMS IOU
+MAX_AGE = 14 # time until missing fish get's new id
+MIN_HITS = 16 # minimum number of frames with a specific fish for it to count
+MIN_LENGTH = 0.3 # minimum fish length, in meters
+IOU_THRES = 0.01 # IOU threshold for tracking
+MIN_TRAVEL = -1 # Minimum distance a track has to travel
+
+class TrackerType(Enum):
+    NONE = 0
+    CONF_BOOST = 1
+    BYTETRACK = 2
+
+class InferenceConfig:
+    def __init__(self, 
+                weights=WEIGHTS, conf_thresh=CONF_THRES, nms_iou=NMS_IOU, 
+                min_hits=MIN_HITS, max_age=MAX_AGE, min_length=MIN_LENGTH, min_travel=MIN_TRAVEL):
+        self.weights = weights
+        self.conf_thresh = conf_thresh
+        self.nms_iou = nms_iou
+        self.min_hits = min_hits
+        self.max_age = max_age
+        self.min_length = min_length
+        self.min_travel = min_travel
+
+        self.associative_tracker = TrackerType.NONE
+        self.boost_power = 1
+        self.boost_decay = 1
+        self.byte_low_conf = 1
+        self.byte_high_conf = 1
+
+    def enable_conf_boost(self, power, decay):
+        self.associative_tracker = TrackerType.CONF_BOOST
+        self.boost_power = power
+        self.boost_decay = decay
+
+    def enable_byte_track(self, low, high):
+        self.associative_tracker = TrackerType.BYTETRACK
+        self.byte_low_conf = low
+        self.byte_high_conf = high
+
+    def to_dict(self):
+        dict = {
+            'weights': self.weights,
+            'nms_iou': self.nms_iou,
+            'min_hits': self.min_hits,
+            'max_age': self.max_age,
+            'min_length': self.min_length,
+            'min_travel': self.min_travel,
+        }
+
+        # Add tracker specific parameters
+        if (self.associative_tracker == TrackerType.BYTETRACK):
+            dict['tracker'] = "ByteTrack"
+            dict['byte_low_conf'] = self.byte_low_conf
+            dict['byte_high_conf'] = self.byte_high_conf
+        elif (self.associative_tracker == TrackerType.CONF_BOOST):
+            dict['tracker'] = "Confidence Boost"
+            dict['conf_thresh'] = self.conf_thresh
+            dict['boost_power'] = self.boost_power
+            dict['boost_decay'] = self.boost_decay
+        elif (self.associative_tracker == TrackerType.NONE):
+            dict['tracker'] = "None"
+            dict['conf_thresh'] = self.conf_thresh
+
+        return dict

app.py

@@ -13,6 +13,7 @@ from gradio_scripts.upload_ui import Upload_Gradio, models
 from gradio_scripts.result_ui import Result_Gradio, update_result, table_headers, info_headers, js_update_tab_labels
 from dataloader import create_dataloader_aris
 from aris import BEAM_WIDTH_DIR
+from InferenceConfig import InferenceConfig
 
 WEBAPP_VERSION = "1.0"
 
@@ -23,7 +24,7 @@ state = {
     'total': 1,
     'annotation_index': -1,
     'frame_index': 0,
-    'hyperparams': {}
+    'config': None
 }
 result = {}
 
@@ -36,26 +37,25 @@ def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_ag
     state['files'] = file_list
     state['total'] = len(file_list)
     state['version'] = WEBAPP_VERSION
-    state['hyperparams'] = {
-        'model': models[model_id] if model_id in models else models['master'],
-        'conf_thresh': conf_thresh,
-        'iou_thresh': iou_thresh,
-        'min_hits': min_hits,
-        'max_age': max_age,
-        'min_length': min_length,
-        'min_travel': min_travel,
-        'associative_tracker': associative_tracker,
-    }
+    state['config'] = InferenceConfig(
+        weights = models[model_id] if model_id in models else models['master'],
+        conf_thresh = conf_thresh,
+        nms_iou = iou_thresh,
+        min_hits = min_hits,
+        max_age = max_age,
+        min_length = min_length,
+        min_travel = min_travel,
+    )
+
+    # Enable tracker if specified
     if (associative_tracker == "Confidence Boost"):
-        state['hyperparams']['boost_power'] = boost_power
-        state['hyperparams']['boost_decay'] = boost_decay
+        state['config'].enable_conf_boost(boost_power, boost_decay)
     elif (associative_tracker == "ByteTrack"):
-        state['hyperparams']['byte_low_conf'] = byte_low_conf
-        state['hyperparams']['byte_high_conf'] = byte_high_conf
+        state['config'].enable_byte_track(byte_low_conf, byte_high_conf)
 
     print(" ")
     print("Running with:")
-    print(state['hyperparams'])
+    print(state['config'].to_dict())
     print(" ")
 
     # Update loading_space to start inference on first file
@@ -169,7 +169,7 @@ def infer_next(_, progress=gr.Progress()):
     # Do inference
     json_result, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(
         file_path, 
-        hyperparams = state['hyperparams'],
+        config = state['config'],
         gradio_progress = set_progress
     )
     

aris.py

@@ -395,6 +395,7 @@ def add_metadata_to_result(aris_fp, json_data, beam_width_dir=BEAM_WIDTH_DIR):
         fish_entry['FRAME_NUM'] = entry['frame_num']
         fish_entry['START_FRAME'] = entry['start_frame_index']
         fish_entry['END_FRAME'] = entry['end_frame_index']
+        fish_entry['NBR_FRAMES'] = entry['end_frame_index'] + 1 - entry['start_frame_index']
         fish_entry['TRAVEL'] = entry['travel_dist']
         fish_entry['DIR'] = upstream_motion_map[entry['direction']]
         fish_entry['R'] = bin_num * pixel_meter_size + frame.windowstart

inference.py

@@ -17,6 +17,7 @@ from lib.yolov5.utils.metrics import box_iou
 import torch
 import torchvision
 
+from InferenceConfig import InferenceConfig, TrackerType
 from lib.fish_eye.tracker import Tracker
 from lib.fish_eye.associative import Associate
 
@@ -50,18 +51,9 @@ def norm(bbox, w, h):
     bb[3] /= h
     return bb
 
-def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None, hyperparams={}):
-
-    # Load hyperparameters
-    if 'model' not in hyperparams: hyperparams['model'] = WEIGHTS
-    if 'iou_thresh' not in hyperparams: hyperparams['iou_thresh'] = NMS_IOU
-    if 'min_hits' not in hyperparams: hyperparams['min_hits'] = MIN_HITS
-    if 'max_age' not in hyperparams: hyperparams['max_age'] = MAX_AGE
-    if 'min_length' not in hyperparams: hyperparams['min_length'] = MIN_LENGTH
-    if 'min_travel' not in hyperparams: hyperparams['min_travel'] = MIN_TRAVEL
-    if 'associative_tracker' not in hyperparams: hyperparams['associative_tracker'] = "None"
+def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None, config=InferenceConfig()):
     
-    model, device = setup_model(hyperparams['model'])
+    model, device = setup_model(config.weights)
 
     load = False
     save = False
@@ -89,43 +81,38 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
         return
 
 
-    outputs = do_suppression(inference, conf_thres=hyperparams['conf_thresh'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+    outputs = do_suppression(inference, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
 
-    if hyperparams['associative_tracker'] == "ByteTrack":
-        if 'byte_low_conf' not in hyperparams: hyperparams['byte_low_conf'] = 0.1
-        if 'byte_high_conf' not in hyperparams: hyperparams['byte_high_conf'] = 0.3
+    if config.associative_tracker == TrackerType.BYTETRACK:
 
-        low_outputs = do_suppression(inference, conf_thres=hyperparams['byte_low_conf'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+        low_outputs = do_suppression(inference, conf_thres=config.byte_low_conf, iou_thres=config.nms_iou, gp=gp)
         low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, gp=gp)
 
-        high_outputs = do_suppression(inference, conf_thres=hyperparams['byte_high_conf'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+        high_outputs = do_suppression(inference, conf_thres=config.byte_high_conf, iou_thres=config.nms_iou, gp=gp)
         high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, gp=gp)
 
         results = do_associative_tracking(
             low_preds, high_preds, image_meter_width, image_meter_height, 
-            reverse=False, min_length=hyperparams['min_length'], min_travel=hyperparams['min_travel'], 
-            max_age=hyperparams['max_age'], min_hits=hyperparams['min_hits'], 
+            reverse=False, min_length=config.min_length, min_travel=config.min_travel, 
+            max_age=config.max_age, min_hits=config.min_hits, 
             gp=gp)
     else: 
 
-        if 'conf_thresh' not in hyperparams: hyperparams['conf_thresh'] = CONF_THRES
 
-        outputs = do_suppression(inference, conf_thres=hyperparams['conf_thresh'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+        outputs = do_suppression(inference, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
 
-        if hyperparams['associative_tracker'] == "Confidence Boost":
-            if 'boost_power' not in hyperparams: hyperparams['boost_power'] = 1
-            if 'boost_decay' not in hyperparams: hyperparams['boost_decay'] = 1
+        if config.associative_tracker == TrackerType.CONF_BOOST:
 
-            do_confidence_boost(inference, outputs, boost_power=hyperparams['boost_power'], boost_decay=hyperparams['boost_decay'], gp=gp)
+            do_confidence_boost(inference, outputs, boost_power=config.boost_power, boost_decay=config.boost_decay, gp=gp)
 
-            outputs = do_suppression(inference, conf_thres=hyperparams['conf_thresh'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+            outputs = do_suppression(inference, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, gp=gp)
 
         all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, gp=gp)
 
         results = do_tracking(
             all_preds, image_meter_width, image_meter_height, 
-            min_length=hyperparams['min_length'], min_travel=hyperparams['min_travel'], 
-            max_age=hyperparams['max_age'], iou_thres=hyperparams['iou_threshold'], min_hits=hyperparams['min_hits'], 
+            min_length=config.min_length, min_travel=config.min_travel, 
+            max_age=config.max_age, min_hits=config.min_hits, 
             gp=gp)
  
     return results

lib/fish_eye/tracker.py

@@ -123,6 +123,7 @@ class Tracker:
 
             fish_entry['travel_dist'] = Tracker.get_travel_distance(start_bbox, end_bbox, json_data['image_meter_width'], json_data['image_meter_height'])
             
+            print(boxes[0])
             fish_entry['start_frame_index'] = boxes[0][1]
             fish_entry['end_frame_index'] = boxes[-1][1]
             fish_entry['color'] = Tracker.selectColor(track_id)

main.py

@@ -7,7 +7,7 @@ from dataloader import create_dataloader_aris
 from inference import do_full_inference, json_dump_round_float
 from visualizer import generate_video_batches
 
-def predict_task(filepath, hyperparams, gradio_progress=None):
+def predict_task(filepath, config, gradio_progress=None):
     """
     Main processing task to be run in gradio
         - Writes aris frames to dirname(filepath)/frames/{i}.jpg
@@ -45,12 +45,12 @@ def predict_task(filepath, hyperparams, gradio_progress=None):
     frame_rate = dataset.didson.info['framerate']
     
     # run detection + tracking
-    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, hyperparams=hyperparams)
+    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, config=config)
 
     # re-index results if desired - this should be done before writing the file
     results = prep_for_mm(results)
     results = add_metadata_to_result(filepath, results)
-    results['metadata']['hyperparameters'] = hyperparams
+    results['metadata']['hyperparameters'] = config.to_dict()
     
     # write output to disk
     json_dump_round_float(results, results_filepath)