New Elwha Model

app.py

@@ -21,13 +21,13 @@ state = {
     'total': 1,
     'annotation_index': -1,
     'frame_index': 0,
-    'model': None
+    'hyperparams': {}
 }
 result = {}
 
 
 # Called when an Aris file is uploaded for inference
-def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_age):
+def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_age, use_associative):
 
     # Reset Result
     reset_state(result, state)
@@ -39,6 +39,7 @@ def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_ag
         'iou_thresh': iou_thresh,
         'min_hits': min_hits,
         'max_age': max_age,
+        'use_associative_tracking': use_associative,
     }
 
     print(" ")
@@ -138,7 +139,7 @@ def infer_next(_, progress=gr.Progress()):
     file_name = file_info[0].split("/")[-1]
     bytes = file_info[1]
     valid, file_path, dir_name = save_data(bytes, file_name)
-    
+
     print("Directory: ", dir_name)
     print("Aris input: ", file_path)
     print(" ")
@@ -153,17 +154,11 @@ def infer_next(_, progress=gr.Progress()):
     # Send uploaded file to AWS 
     upload_file(file_path, "fishcounting", "webapp_uploads/" + file_name)
 
-    hyperparams = state['hyperparams']
-
     # Do inference
     json_result, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(
         file_path, 
-        weights = hyperparams['model'], 
-        conf_thresh = hyperparams['conf_thresh'], 
-        iou_thresh = hyperparams['iou_thresh'], 
-        min_hits = hyperparams['min_hits'], 
-        max_age = hyperparams['max_age'], 
-        gradio_progress=set_progress
+        hyperparams = state['hyperparams'],
+        gradio_progress = set_progress
     )
     
     # Store result for that file

gradio_scripts/upload_ui.py

@@ -4,7 +4,8 @@ from gradio_scripts.file_reader import File
 
 models = {
     'master': 'models/v5m_896_300best.pt',
-    'elwha': 'models/YsEE20.pt'
+    'elwha': 'models/YsEE20.pt',
+    'elwha_train': 'models/YsEKtE20.pt',
 }
 
 def Upload_Gradio(gradio_components):
@@ -29,6 +30,13 @@ def Upload_Gradio(gradio_components):
                     settings.append(gr.Slider(0, 100, value=16, label="Min Hits", info="Minimum number of frames a fish has to appear in to count"))
                     settings.append(gr.Slider(0, 100, value=14, label="Max Age", info="Max age of occlusion before track is split"))
 
+                with gr.Row():
+                    gr.Markdown("Associative Tracking")
+                    settings.append(gr.Checkbox(value=False, label="Enabled"))
+                with gr.Row():
+                    settings.append(gr.Slider(0, 100, value=16, label="Min Hits", info="Minimum number of frames a fish has to appear in to count"))
+                    settings.append(gr.Slider(0, 100, value=14, label="Max Age", info="Max age of occlusion before track is split"))
+
                 gradio_components['hyperparams'] = settings
 
             #Input field for aris submission

inference.py

@@ -48,9 +48,18 @@ def norm(bbox, w, h):
     bb[3] /= h
     return bb
 
-def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None, weights=WEIGHTS, conf_thresh=CONF_THRES, nms_iou=NMS_IOU, min_hits=MIN_HITS, max_age=MAX_AGE):
+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 'conf_thresh' not in hyperparams: hyperparams['conf_tresh'] = CONF_THRES
+    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 'use_associative_tracking' not in hyperparams: hyperparams['use_associative_tracking'] = False
+    if 'AT_decay' not in hyperparams: hyperparams['AT_decay'] = MIN_HITS
     
-    model, device = setup_model(weights)
+    model, device = setup_model(hyperparams['model'])
 
     load = False
     save = False
@@ -78,15 +87,17 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
         return
 
 
-    outputs = do_suppression(inference, conf_thres=conf_thresh, iou_thres=nms_iou, gp=gp)
+    outputs = do_suppression(inference, conf_thres=hyperparams['conf_tresh'], iou_thres=hyperparams['iou_tresh'], gp=gp)
 
-    #do_confidence_boost(inference, outputs, gp=gp)
+    if hyperparams['use_associative_tracking']:
+        
+        do_confidence_boost(inference, outputs, gp=gp)
 
-    #new_outputs = do_suppression(inference, conf_thres=conf_thresh, iou_thres=nms_iou, gp=gp)
+        outputs = do_suppression(inference, conf_thres=hyperparams['conf_tresh'], iou_thres=hyperparams['iou_tresh'], 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_hits=min_hits, max_age=max_age, gp=gp)
+    results = do_tracking(all_preds, image_meter_width, image_meter_height, min_hits=hyperparams['min_hits'], max_age=hyperparams['max_age'], gp=gp)
 
     return results
     

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, weights, conf_thresh, iou_thresh, min_hits, max_age, gradio_progress=None):
+def predict_task(filepath, hyperparams, gradio_progress=None):
     """
     Main processing task to be run in gradio
         - Writes aris frames to dirname(filepath)/frames/{i}.jpg
@@ -45,18 +45,12 @@ def predict_task(filepath, weights, conf_thresh, iou_thresh, min_hits, max_age,
     frame_rate = dataset.didson.info['framerate']
     
     # run detection + tracking
-    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, weights=weights, conf_thresh=conf_thresh, nms_iou=iou_thresh, min_hits=min_hits, max_age=max_age)
+    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, hyperparams=hyperparams)
 
     # 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'] = {
-        'model': weights,
-        'conf_thresh': conf_thresh,
-        'iou_thresh': iou_thresh,
-        'min_hits': min_hits,
-        'max_age': max_age
-    }
+    results['metadata']['hyperparameters'] = hyperparams
     
     # write output to disk
     json_dump_round_float(results, results_filepath)