Added Elwha test model

app.py

@@ -9,7 +9,7 @@ from gradio_scripts.annotation_handler import init_frames
 import json
 from zipfile import ZipFile
 import os
-from gradio_scripts.upload_ui import Upload_Gradio
+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
@@ -20,18 +20,23 @@ state = {
     'index': 1,
     'total': 1,
     'annotation_index': -1,
-    'frame_index': 0
+    'frame_index': 0,
+    'model': None
 }
 result = {}
 
 
 # Called when an Aris file is uploaded for inference
-def on_aris_input(file_list):
+def on_aris_input(file_list, model_id):
+
+    print(model_id)
+    print(models[model_id] if model_id in models else models['master'])
 
     # 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']
 
     # Update loading_space to start inference on first file
     return {
@@ -139,7 +144,7 @@ def infer_next(_, progress=gr.Progress()):
     upload_file(file_path, "fishcounting", "webapp_uploads/" + file_name)
 
     # Do inference
-    json_result, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(file_path, gradio_progress=set_progress)
+    json_result, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(file_path, weights=state['model'], gradio_progress=set_progress)
     
     # Store result for that file
     result['json_result'].append(json_result)
@@ -365,7 +370,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, inference_comps)
+    input.upload(on_aris_input, [input, components['model_select']], 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/annotation_handler.py

@@ -39,7 +39,7 @@ def init_frames(dataset, preds, index, gp=None):
             if gp: gp((index + i)/len(preds['frames']), "Extracting Frames")
 
             # Extract frames
-            img_raw = dataset.didson.load_frames(start_frame=i, end_frame=i+1)[0]
+            img_raw = dataset.didson.load_frames(start_frame=index+i, end_frame=index+i+1)[0]
             image = cv2.resize(cv2.cvtColor(img_raw, cv2.COLOR_GRAY2BGR), (w, h))
             #cv2.imwrite("annotation_frame_dir/" + str(i) + ".jpg", image)
             retval, buffer = cv2.imencode('.jpg', image)

gradio_scripts/upload_ui.py

@@ -2,6 +2,11 @@ import gradio as gr
 from gradio_scripts.file_reader import File
 
 
+models = {
+    'master': 'models/v5m_896_300best.pt',
+    'elwha': 'models/YsEE20.pt'
+}
+
 def Upload_Gradio(gradio_components):
     with gr.Tabs():
 
@@ -10,6 +15,8 @@ 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()))
+
             #Input field for aris submission
             gradio_components['input'] = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")
 

models/YsEE20.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d661200690bc0e9d3e48b907fcf0a9fa9165b3e92861e275b21cc962c5970262
+size 56791439