Hyperparameter settings

app.py

@@ -27,16 +27,24 @@ result = {}
 
 
 # Called when an Aris file is uploaded for inference
-def on_aris_input(file_list, model_id):
-
-    print(model_id)
-    print(models[model_id] if model_id in models else models['master'])
+def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_age):
 
     # Reset Result
     reset_state(result, state)
     state['files'] = file_list
     state['total'] = len(file_list)
-    state['model'] = models[model_id] if model_id in models else models['master']
+    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,
+    }
+
+    print(" ")
+    print("Running with:")
+    print(state['hyperparams'])
+    print(" ")
 
     # Update loading_space to start inference on first file
     return {
@@ -130,8 +138,10 @@ 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(dir_name)
-    print(file_path)
+    
+    print("Directory: ", dir_name)
+    print("Aris input: ", file_path)
+    print(" ")
 
     # Check that the file was valid
     if not valid: 
@@ -143,8 +153,18 @@ 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=state['model'], gradio_progress=set_progress)
+    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
+    )
     
     # Store result for that file
     result['json_result'].append(json_result)
@@ -370,7 +390,7 @@ with demo:
     inference_comps = [inference_handler, master_tabs, components['cancelBtn'], components['skipBtn']]
 
     # When a file is uploaded to the input, tell the inference_handler to start inference
-    input.upload(on_aris_input, [input, components['model_select']], inference_comps)
+    input.upload(on_aris_input, [input] + components['hyperparams'], inference_comps)
 
     # When inference handler updates, tell result_handler to show the new result
     # Also, add inference_handler as the output in order to have it display the progress

gradio_scripts/upload_ui.py

@@ -15,7 +15,21 @@ def Upload_Gradio(gradio_components):
 
             gr.HTML("<p align='center' style='font-size: large;font-style: italic;'>Submit an .aris file to analyze result.</p>")
             
-            gradio_components['model_select'] = gr.Dropdown(value="master", choices=list(models.keys()))
+            with gr.Accordion("Advanced Settings", open=False):
+                settings = []
+                settings.append(gr.Dropdown(label="Model", value="master", choices=list(models.keys())))
+
+                gr.Markdown("Detection Parameters")
+                with gr.Row():
+                    settings.append(gr.Slider(0, 1, value=0.05, label="Confidence Threshold", info="Confidence cutoff for detection boxes"))
+                    settings.append(gr.Slider(0, 1, value=0.2, label="NMS IoU", info="IoU threshold for non-max suppression"))
+
+                gr.Markdown("Tracking Parameters")
+                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
             gradio_components['input'] = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")

inference.py

@@ -48,7 +48,7 @@ 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):
+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):
     
     model, device = setup_model(weights)
 
@@ -78,15 +78,15 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
         return
 
 
-    outputs = do_suppression(inference, gp=gp)
+    outputs = do_suppression(inference, conf_thres=conf_thresh, iou_thres=nms_iou, gp=gp)
 
     #do_confidence_boost(inference, outputs, gp=gp)
 
-    #new_outputs = do_suppression(inference, gp=gp)
+    #new_outputs = do_suppression(inference, conf_thres=conf_thresh, iou_thres=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, gp=gp)
+    results = do_tracking(all_preds, image_meter_width, image_meter_height, min_hits=min_hits, max_age=max_age, gp=gp)
 
     return results
     

main.py

@@ -7,9 +7,7 @@ from dataloader import create_dataloader_aris
 from inference import do_full_inference, json_dump_round_float
 from visualizer import generate_video_batches
 
-WEIGHTS = 'models/v5m_896_300best.pt'
-
-def predict_task(filepath, weights=WEIGHTS, gradio_progress=None):
+def predict_task(filepath, weights, conf_thresh, iou_thresh, min_hits, max_age, gradio_progress=None):
     """
     Main processing task to be run in gradio
         - Writes aris frames to dirname(filepath)/frames/{i}.jpg
@@ -25,7 +23,6 @@ def predict_task(filepath, weights=WEIGHTS, gradio_progress=None):
     if (gradio_progress): gradio_progress(0, "In task...")
     print("Cuda available in task?", torch.cuda.is_available())
     
-    print(filepath)
     dirname = os.path.dirname(filepath)
     filename = os.path.basename(filepath).replace(".aris","").replace(".ddf","")
     results_filepath = os.path.join(dirname, f"{filename}_results.json")
@@ -48,7 +45,7 @@ def predict_task(filepath, weights=WEIGHTS, 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, weights=weights)
+    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)
 
     # re-index results if desired - this should be done before writing the file
     results = prep_for_mm(results)

visualizer.py

@@ -110,7 +110,7 @@ def get_video_frames(frames, preds, frame_rate, image_meter_width=None, image_me
             cv2.putText(frame, f'Left count: {clip_pr_counts[FONT_THICKNESS]}',  (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*2), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
             cv2.putText(frame, f'Other fish: {clip_pr_counts[2]}',  (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*1), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
     #         cv2.putText(frame, f'Upstream: {preds["upstream_direction"]}', (0, h-1-LINE_HEIGHT*1), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
-            cv2.putText(frame, f'Frame: {i}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*0), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+            cv2.putText(frame, f'Frame: {start_frame+i}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*0), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
 
             vid_frames.append(frame)