Compatability with backend

.gitignore

@@ -6,6 +6,7 @@ static/tmp.jpg
 redis-stable/*
 user_data/*
 models/*
+tmp/*
 *.pyc
 
 .ipynb_checkpoints
@@ -14,4 +15,5 @@ models/*
 *.aris
 *.log
 *.pdf
+*.mp4
 *.DS_STORE

app.py

@@ -1,23 +1,22 @@
 import gradio as gr
-from uploader import save_data_to_dir, create_data_dir, save_data
-from main import predict_task
-from gradio_scripts.state_handler import reset_state
 import numpy as np
-from gradio_scripts.aws_handler import upload_file
-from aris import create_metadata_table
-from gradio_scripts.annotation_handler import init_frames
 import json
 from zipfile import ZipFile
 import os
-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"
+from backend.dataloader import create_dataloader_aris
+from backend.aws_handler import ping_server
+from backend.predict import predict_task
+from backend.uploader import save_data_to_dir, create_data_dir, save_data
+from backend.InferenceConfig import InferenceConfig
 
-enable_annotation_editor = False
+from frontend.upload_ui import Upload_Gradio, models
+from frontend.result_ui import Result_Gradio, update_result, create_metadata_table, table_headers, info_headers
+from frontend.annotation_handler import load_annotation, prepare_annotation, js_store_frame_info, annotation_css
+from frontend.state_handler import reset_state
+
+
+WEBAPP_VERSION = "Advanced 1.0"
 
 #Initialize State & Result
 state = {
@@ -28,11 +27,13 @@ state = {
     'frame_index': 0,
     'outputs': [],
     'config': None,
+    'enable_annotation_editor': False
 }
 result = {}
+components = {}
 
-
-# Called when an Aris file is uploaded for inference
+# -------------------------------------------- ----- UPLOAD ARIS FILE ------------------------------------------------------
+# Called when an Aris file is uploaded for inference - calls infer_next
 def on_aris_input(
         file_list, 
         model_id, 
@@ -43,6 +44,9 @@ def on_aris_input(
         output_formats
     ):
 
+    if isinstance(file_list, tuple):
+        file_list = [file_list]
+
     print(output_formats)
 
     # Reset Result
@@ -82,77 +86,6 @@ def on_aris_input(
         master_tabs: gr.update(selected=1)
     }
 
-# Called when a result zip file is uploaded for result review
-def on_result_upload():
-    return {
-        master_tabs: gr.update(selected=1),
-        result_uploader: gr.update(value=str(np.random.rand()))
-    }
-
-def on_result_upload_finish(zip_list, aris_list):
-
-    if (zip_list == None): 
-        zip_list = [("static/example/example_result.zip", None)]
-        aris_path = "static/example/input_file.aris"
-        aris_list = [(aris_path, bytearray(open(aris_path, 'rb').read()))]
-    
-
-    reset_state(result, state)
-    state['version'] = WEBAPP_VERSION
-    state['outputs'] = ["Annotated Video", "Manual Marking", "PDF"]
-
-    component_updates = { 
-        tab_labeler: gr.update(value = len(zip_list))
-    }
-
-    for i in range(len(zip_list)):
-            
-        # Create dir to unzip files
-        dir_name = create_data_dir(str(i))
-
-        # Check aris input 
-        if (aris_list):
-            aris_info = aris_list[i]
-            file_name = aris_info[0].split("/")[-1]
-            bytes = aris_info[1]
-            valid, input_path, dir_name = save_data_to_dir(bytes, file_name, dir_name)
-        else:
-            input_path = None
-
-        # Unzip result
-        zip_info = zip_list[i]
-        zip_name = zip_info[0]
-        print(zip_name)
-        with ZipFile(zip_name) as zip_file:
-            ZipFile.extractall(zip_file, path=dir_name)
-        unzipped = os.listdir(dir_name)
-        print(unzipped)
-
-        for file in unzipped:
-            if (file.endswith("_results.mp4")):
-                result["path_video"].append(os.path.join(dir_name, file))
-            elif (file.endswith("_results.json")):
-                result["path_json"].append(os.path.join(dir_name, file))
-            elif (file.endswith("_marking.txt")):
-                result["path_marking"].append(os.path.join(dir_name, file))
-        
-        result["aris_input"].append(input_path)
-        with open(result['path_json'][-1]) as f:
-            json_result = json.load(f)
-            result['json_result'].append(json_result)
-        fish_table, fish_info = create_metadata_table(json_result, table_headers, info_headers)
-        result["fish_table"].append(fish_table)
-        result["fish_info"].append(fish_info)
-
-        update = update_result(i, state, result, inference_handler)
-
-        for key in update.keys():
-            component_updates[key] = update[key]
-    
-    component_updates.pop(inference_handler)
-    return component_updates
-
-
 # Iterative function that performs inference on the next file in line
 def infer_next(_, progress=gr.Progress()):
 
@@ -186,7 +119,10 @@ def infer_next(_, progress=gr.Progress()):
         }
 
     # Send uploaded file to AWS 
-    upload_file(file_path, "fishcounting", "webapp_uploads/" + file_name)
+    ping_server(file_name, state)
+    #upload_file(file_path, "fishcounting", "webapp_uploads/files/" + file_name)
+
+    #crop_clip(file_path, 65)
 
     # Do inference
     json_result, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(
@@ -195,10 +131,14 @@ def infer_next(_, progress=gr.Progress()):
         output_formats = state['outputs'],
         gradio_progress = set_progress
     )
+
+    # prepare dummy dataloader for visualizations
+    _, dataset = create_dataloader_aris(file_path, num_frames_bg_subtract=0)
     
     # Store result for that file
     result['json_result'].append(json_result)
     result['aris_input'].append(file_path)
+    result['datasets'].append(dataset)
     result["path_video"].append(video_filepath)
     result["path_zip"].append(zip_filepath)
     result["path_json"].append(json_filepath)
@@ -218,12 +158,7 @@ def infer_next(_, progress=gr.Progress()):
         inference_handler: gr.update()
     }
 
-# Show result
-def on_result_ready():
-    # Update result tab for last file
-    i = state["index"] - 1
-    return update_result(i, state, result, inference_handler)
-
+# Cancel inference
 def cancel_inference():
     return {
         master_tabs: gr.update(selected=0),
@@ -231,117 +166,118 @@ def cancel_inference():
         components['cancel_btn']: gr.update(visible=False)
     }
 
+# Show result
+def on_result_ready():
+    # Update result tab for last file
+    i = state["index"] - 1
+    return update_result(i, state, result, inference_handler)
 
-# Request loading of animation editor
-def prepare_annotation(index):
 
-    state['annotation_index'] = index
-    state['frame_index'] = 0
 
-    if result["aris_input"][index]:
-        return {
-            annotation_progress: gr.update(value="<p id='annotation_info' style='display:none'>[]</p><!--" + str(np.random.rand()) + "-->", visible=True),
-            master_tabs: gr.update(selected=2)
-        }
+# ------------------------------------------------- UPLOAD RESULT FILE -----------------------------------------------------
+# Called when result file is uploaded for review
+def on_result_upload():
     return {
-        annotation_progress: gr.update(),
-        master_tabs: gr.update()
+        master_tabs: gr.update(selected=1),
+        result_uploader: gr.update(value=str(np.random.rand()))
     }
 
-annotation_info = None
-annotation_dataset = None
-# annotation_progress.change
-def load_annotation(_, progress=gr.Progress()):
-    global annotation_info, annotation_dataset
-
-    # Get result index
-    result_index = state['annotation_index']
+# Called when result upload is finished processing
+def on_result_upload_finish(zip_list, aris_list):
 
-    set_progress = lambda pct, msg: progress(pct, desc=msg)
+    if (zip_list == None): 
+        zip_list = [("static/example/example_result.zip", None)]
+        aris_path = "static/example/input_file.aris"
+        aris_list = [(aris_path, bytearray(open(aris_path, 'rb').read()))]
+    
 
-    if state['frame_index'] == 0:
-        if set_progress: set_progress(0, "Loading Frames")
-        dataloader, annotation_dataset = create_dataloader_aris(result["aris_input"][result_index], BEAM_WIDTH_DIR, None)
+    reset_state(result, state)
+    state['version'] = WEBAPP_VERSION
+    state['outputs'] = ["Generate Annotated Video", "Generate Manual Marking", "Generate PDF"]
 
-    # Check that frames remain to be loaded
-    if state['frame_index'] < len(result['json_result'][result_index]['frames']):
+    component_updates = { 
+        tab_labeler: gr.update(value = len(zip_list))
+    }
 
-        # load frames and annotation
-        annotation_info, state['frame_index'] = init_frames(annotation_dataset, result['json_result'][result_index], state['frame_index'], gp=set_progress)
-    
-        # save as html element
-        annotation_content = "<p id='annotation_info' style='display:none'>" + json.dumps(annotation_info) + "</p>"
+    for i in range(len(zip_list)):
+            
+        # Create dir to unzip files
+        dir_name = create_data_dir(str(i))
 
-        return {
-            annotation_editor: gr.update(), 
-            annotation_progress: gr.update(value=annotation_content)
-        }
+        # Check aris input 
+        if (aris_list):
+            aris_info = aris_list[i]
+            file_name = aris_info[0].split("/")[-1]
+            bytes = aris_info[1]
+            valid, input_path, dir_name = save_data_to_dir(bytes, file_name, dir_name)
+            _, dataset = create_dataloader_aris(input_path, num_frames_bg_subtract=0)
+        else:
+            input_path = None
+            dataset = None
 
-    # If complete, start annotation editor    
+        # Unzip result
+        zip_info = zip_list[i]
+        zip_name = zip_info[0]
+        print(zip_name)
+        with ZipFile(zip_name) as zip_file:
+            ZipFile.extractall(zip_file, path=dir_name)
+        unzipped = os.listdir(dir_name)
+        print(unzipped)
 
-    annotation_html = ""
+        for file in unzipped:
+            if (file.endswith("_results.mp4")):
+                result["path_video"].append(os.path.join(dir_name, file))
+            elif (file.endswith("_results.json")):
+                result["path_json"].append(os.path.join(dir_name, file))
+            elif (file.endswith("_marking.txt")):
+                result["path_marking"].append(os.path.join(dir_name, file))
+        
+        result["aris_input"].append(input_path)
+        result["datasets"].append(dataset)
+        with open(result['path_json'][-1]) as f:
+            json_result = json.load(f)
+            result['json_result'].append(json_result)
+        fish_table, fish_info = create_metadata_table(json_result, table_headers, info_headers)
+        result["fish_table"].append(fish_table)
+        result["fish_info"].append(fish_info)
 
-    # Header
-    annotation_html += "<div id='annotation_header'>"
-    annotation_html += "     <h1 id='annotation_frame_nbr'>Frame 0/100</h1>"
-    annotation_html += "     <p id='annotation_edited'>(edited)</p>"
-    annotation_html += "</div>"
+        update = update_result(i, state, result, inference_handler)
 
-    # Annotation Body
-    annotation_html += "<div style='display:flex'>"
-    annotation_html += "     <canvas id='canvas' style='width:50%' onmousedown='mouse_down(event)' onmousemove='mouse_move(event)' onmouseup='mouse_up()' onmouseleave='mouse_up()'></canvas>"
-    annotation_html += "     <div id='annotation_display' style='width:50%'></div>"
-    annotation_html += "</div>"
+        for key in update.keys():
+            component_updates[key] = update[key]
     
-    # Dummy objects
-    annotation_html += "<img id='annotation_img' onload='draw()' style='display:none'></img>"
-    annotation_html += "<!--" + str(np.random.rand()) + "-->"
+    component_updates.pop(inference_handler)
+    return component_updates
 
-    return {
-        annotation_editor: gr.update(value=annotation_html, visible=True),
-        annotation_progress: gr.update(visible=False)
-    }
 
 
-components = {}
+# ------------------------------------------------- ANNOTATION EDITOR -----------------------------------------------------
+def on_annotation_open(result_index):
+    return prepare_annotation(state, result, result_index)
+
+def annotate_next(_, progress=gr.Progress()):
+    return load_annotation(state, result, progress)
+
+
 
-demo = gr.Blocks()
-with demo:
-    with gr.Blocks() as inner_body:
+# -------------------------------------------------- GRADIO ARCHITECTURE ----------------------------------------------------
+with gr.Blocks() as demo:
+    with gr.Blocks():
 
         # Title of page + style
         gr.HTML(
             """
-            <h1 align="center" style="font-size:xxx-large">Caltech Fisheye</h1>
+            <h1 align="center" style="font-size:xxx-large">Caltech Fisheye - Advanced</h1>
             <style>
+                /* Disable header of metadata list in result */
                 #marking_json thead {
                     display: none !important;
                 }
+                /* Color of selected tab */
                 .selected.svelte-kqij2n {
                     background: linear-gradient(180deg, #66eecb47, transparent);
                 }
-                #annotation_frame_nbr {
-                    left: calc(50% - 100px);
-                    position: absolute;
-                    width: 200px;
-                    text-align: center;
-                    font-size: x-large;
-                }
-                #annotation_header {
-                    height: 40px;
-                }
-                #annotation_frame_nbr {
-                    left: calc(50% - 100px);
-                    position: absolute;
-                    width: 200px;
-                    text-align: center;
-                    font-size: x-large;
-                }
-                #annotation_edited {
-                    right: 0px;
-                    position: absolute;
-                    margin-top: 5px;
-                }
+                """ + annotation_css + """
             </style>
             <style id="tab_style"></style>
         """
@@ -360,7 +296,6 @@ with demo:
             # Master Tab for result visualization
             with gr.Tab("Result", id=1):
 
-
                 # Define annotation progress bar for event listeres, but unrender since it will be displayed later on
                 result_uploader = gr.HTML("", visible=False)
                 components['result_uploader'] = result_uploader
@@ -369,10 +304,10 @@ with demo:
                 components['annotation_progress'] = annotation_progress
                 
                 # Draw the gradio components related to visualzing result
-                vis_components = Result_Gradio(prepare_annotation, components)
+                visualization_components = Result_Gradio(on_annotation_open, components, state)
                 
             # Master Tab for annotation editing
-            if enable_annotation_editor:
+            if state['enable_annotation_editor']:
                 with gr.Tab("Annotation Editor", id=2):
 
                     # Draw the annotation loading bar here
@@ -381,21 +316,13 @@ with demo:
                     # Add annotation editor component
                     annotation_editor = gr.HTML("", visible=False)
 
-                    # Event listener for opening annotation 
-                    annotation_progress.change(load_annotation, annotation_progress, [annotation_editor, annotation_progress], _js="""
-                        () => {
-                            info_string = document.getElementById("annotation_info").innerHTML;
-                            info = JSON.parse(info_string);
-                            console.log(info)
-                            if (info.length == 0) {
-                                window.annotation_info = [];
-                                return false;
-                            } 
-                            window.annotation_info = window.annotation_info.concat(info)
-                            console.log(window.annotation_info)
-                            return true;
-                        }
-                    """)
+                    # Event listener for batch loading of annotation frames 
+                    annotation_progress.change(
+                        annotate_next, 
+                        annotation_progress, 
+                        [annotation_editor, annotation_progress], 
+                        _js=js_store_frame_info
+                    )
                     
                     # Event listener for running javascript defined in 'annotation_editor.js'
                     # show_annotation
@@ -417,19 +344,18 @@ with demo:
     inference_handler = components['inference_handler']
     result_handler = components['result_handler']
     tab_labeler = components['tab_labeler']
-
-
     inference_comps = [inference_handler, master_tabs, components['cancel_btn']]
 
     # When a file is uploaded to the input, tell the inference_handler to start inference
     input.upload(on_aris_input, [input] + components['hyperparams'], inference_comps)
+    components['inference_btn'].click(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
     inference_event = inference_handler.change(infer_next, None, [inference_handler, result_handler, tab_labeler])
 
     # Send UI changes based on the new results to the UI_components, and tell the inference_handler to start next inference
-    result_handler.change(on_result_ready, None, vis_components + [inference_handler])
+    result_handler.change(on_result_ready, None, visualization_components + [inference_handler])
 
     # Cancel and skip buttons
     components['cancel_btn'].click(cancel_inference, None, inference_comps, cancels=[inference_event])
@@ -439,11 +365,7 @@ with demo:
     components['result_uploader'].change(
         on_result_upload_finish, 
         [components['result_input'], components['result_aris_input']], 
-        vis_components + [tab_labeler]
+        visualization_components + [tab_labeler]
     )
 
 demo.queue().launch()
-
-on_result_ready()
-
-

InferenceConfig.py → backend/InferenceConfig.py

@@ -47,16 +47,42 @@ class InferenceConfig:
     def enable_sort_track(self):
         self.associative_tracker = TrackerType.NONE
 
-    def enable_conf_boost(self, power, decay):
+    def enable_conf_boost(self, power=2, decay=0.1):
         self.associative_tracker = TrackerType.CONF_BOOST
         self.boost_power = power
         self.boost_decay = decay
 
-    def enable_byte_track(self, low, high):
+    def enable_byte_track(self, low=0.1, high=0.3):
         self.associative_tracker = TrackerType.BYTETRACK
         self.byte_low_conf = low
         self.byte_high_conf = high
 
+    def enable_tracker_from_string(self, associativity):
+        if associativity != "":
+            if (associativity.startswith("boost")):
+                conf = associativity.split(":")
+                if len(conf) == 3:
+                    self.enable_conf_boost(power=float(conf[1]), decay=float(conf[2]))
+                    return True
+                else:
+                    print("INVALID PARAMETERS FOR CONFIDENCE BOOST:", associativity)
+                return False
+                
+            elif (associativity.startswith("bytetrack")):
+                conf = associativity.split(":")
+                if len(conf) == 3:
+                    self.enable_byte_track(low=float(conf[1]), high=float(conf[2]))
+                    return True
+                else:
+                    print("INVALID PARAMETERS FOR BYTETRACK:", associativity)
+                return False
+            else:
+                print("INVALID ASSOCIATIVITY TYPE:", associativity)
+                return False
+        else:
+            self.enable_sort_track()
+            return True
+
     def find_model(self, model_list):
         print("weights", self.weights)
         for model_name in model_list:

aris.py → backend/aris.py

@@ -9,7 +9,6 @@ from PIL import Image
 from tqdm import tqdm
 import datetime
 from decimal import Decimal, ROUND_HALF_UP
-import json
 import pytz
 from copy import deepcopy
 from multiprocessing import Pool
@@ -28,7 +27,12 @@ ImageData = namedtuple('ImageData', [
     'sample_read_rows', 'sample_read_cols', 'image_write_rows', 'image_write_cols'
 ])
 
-def FastARISRead(ARIS_data, start_frame, end_frame):
+def FastARISRead(aris_fp, start_frame, end_frame):
+    ARISdata, aris_frame = pyARIS.DataImport(aris_fp)
+    frames = FastARISExtract(ARISdata, start_frame, end_frame)
+    return frames
+
+def FastARISExtract(ARIS_data, start_frame, end_frame):
     """ Just read in the ARIS frame, and not the other meta data. 
     """
     FrameSize = ARIS_data.SamplesPerChannel*ARIS_data.NumRawBeams
@@ -61,7 +65,7 @@ def get_info(aris_fp, beam_width_dir=BEAM_WIDTH_DIR):
     )
     return pixel_meter_size * xdim, pixel_meter_size * ydim, aris_frame.framerate
 
-def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH_DIR, bg_out_dir=None, num_workers=0):
+def write_frames(aris_fp, out_dir, max_mb=-1, beam_width_dir=BEAM_WIDTH_DIR, bg_out_dir=None, num_workers=0):
     """
     Write all frames from an ARIS file to disk, using our 3-channel format:
     (raw img, blurred & mean subtracted img, optical flow approximation)
@@ -69,7 +73,6 @@ def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH
     Args:
         aris_fp: path to aris file
         out_dir: directory for frame extraction. frames will be named 0.jpg, 1.jpg, ... {n}.jpg
-        cb: a callback function for updating progress
         max_mb: maximum amount of the file to be processed, in megabytes
         beam_width_dir: location of ARIS camera information
         bg_out_dir: where to write the background frame; None disables writing
@@ -82,9 +85,6 @@ def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH
     # Load in the ARIS file
     ARISdata, aris_frame = pyARIS.DataImport(aris_fp)
     
-    if cb:
-        cb(2, msg="Decoding ARIS data...")
-    
     beam_width_data = pyARIS.load_beam_width_data(aris_frame, beam_width_dir=beam_width_dir)[0]
     # What is the meter resolution of the smallest sample?
     min_pixel_size = pyARIS.get_minimum_pixel_meter_size(aris_frame, beam_width_data)
@@ -135,18 +135,18 @@ def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH
     with tqdm(total=(end_frame-start_frame-1), desc="Extracting frames", ncols=0) as pbar:
         # compute info for bg subtraction using first batch
         # TODO: make this a sliding window
-        mean_blurred_frame, mean_normalization_value = write_frame_range(ARISdata, image_data, out_dir, clips[0][0], clips[0][1], None, None, cb, pbar)
+        mean_blurred_frame, mean_normalization_value = write_frame_range(ARISdata, image_data, out_dir, clips[0][0], clips[0][1], None, None, pbar)
     
         # do rest of batches in parallel
         if num_workers > 0:
-            args = [ (ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, cb) for (start, end) in clips[1:] ] # TODO: can't pass pbar to thread
+            args = [ (ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value) for (start, end) in clips[1:] ] # TODO: can't pass pbar to thread
             with Pool(num_workers) as pool:
                 results = [ pool.apply_async(write_frame_range, arg) for arg in args ]
                 results = [ r.get() for r in results ] # need this call to block on thread execution
                 pbar.update(sum([ arg[4] - arg[3] for arg in args ]))
         else:
             for j, (start, end) in enumerate(clips[1:]):
-                write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, cb, pbar)
+                write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, pbar)
             
     if bg_out_dir is not None:
         bg_img = (mean_blurred_frame * 255).astype(np.uint8)
@@ -155,10 +155,10 @@ def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH
         
     return pixel_meter_size * xdim, pixel_meter_size * ydim, aris_frame.framerate
 
-def write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame=None, mean_normalization_value=None, cb=None, pbar=None):
+def write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame=None, mean_normalization_value=None, pbar=None):
     try:
         frames = np.zeros([end-start, image_data.ydim, image_data.xdim], dtype=np.uint8)
-        frames[:, image_data.image_write_rows, image_data.image_write_cols] = FastARISRead(ARISdata, start, end)[:, image_data.sample_read_rows, image_data.sample_read_cols]
+        frames[:, image_data.image_write_rows, image_data.image_write_cols] = FastARISExtract(ARISdata, start, end)[:, image_data.sample_read_rows, image_data.sample_read_cols]
     except:
         print("Error extracting frames from", ARISdata.filename, "during batch", i)
         return
@@ -196,12 +196,12 @@ def write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_fr
         
         if pbar:
             pbar.update(1)
-        if cb:
-            pct = 2 + int( (start+i) / (end_frame - start_frame - 1) * 98)
-            cb(pct, msg=pbar.__str__())
             
     return mean_blurred_frame, mean_normalization_value
 
+
+
+
 def prep_for_mm(json_data):
     """Prepare json results for writing to a manual marking file."""
     json_data = deepcopy(json_data)
@@ -257,7 +257,6 @@ def prep_for_mm(json_data):
         
     return json_data
 
-
 def add_metadata_to_result(aris_fp, json_data, beam_width_dir=BEAM_WIDTH_DIR):
     """
     Return:
@@ -436,49 +435,6 @@ def add_metadata_to_result(aris_fp, json_data, beam_width_dir=BEAM_WIDTH_DIR):
 
     return json_data
 
-def create_metadata_table(result, table_headers, info_headers):
-    if 'metadata' in result:
-        metadata = result['metadata']
-    else:
-        metadata = { 'FISH': [] }
-    
-    # Calculate detection dropout
-    for fish in metadata['FISH']:
-        count = 0
-        for frame in result['frames'][fish['START_FRAME']:fish['END_FRAME']+1]:
-            for ann in frame['fish']:
-                if ann['fish_id'] == fish['TOTAL']:
-                    count += 1
-        fish['DETECTION_DROPOUT'] = 1 - count / (fish['END_FRAME'] + 1 - fish['START_FRAME'])
-
-    # Create fish table
-    table = []
-    for fish in metadata["FISH"]:
-        row = []
-        for header in table_headers:
-            row.append(fish[header])
-        table.append(row)
-    
-    if len(metadata["FISH"]) == 0:
-        row = []
-        for header in table_headers:
-            row.append("-")
-        table.append(row)
-
-    # Create info table
-    info = []
-    for field in info_headers:
-        field_name = "**" + field + "**"
-        if field in metadata:
-            info.append([field_name, str(metadata[field])])
-        else:
-            info.append([field_name, ""])
-    if 'hyperparameters' in metadata:
-        for param_name in metadata['hyperparameters']:
-            info.append(['**' + param_name + '**', str(metadata['hyperparameters'][param_name])])
-    
-    return table, info
-
 def create_manual_marking(results, out_path=None):
     """
     Return:

{gradio_scripts → backend}/aws_handler.py

@@ -2,9 +2,11 @@ import logging
 import boto3
 from botocore.exceptions import ClientError
 import os
+from datetime import datetime
+import torch
 
 
-def upload_file(file_name, bucket, object_name=None):
+def upload_file(file_name, bucket="fishcounting", object_name=None):
     """Upload a file to an S3 bucket
 
     :param file_name: File to upload
@@ -32,4 +34,45 @@ def upload_file(file_name, bucket, object_name=None):
     except ClientError as e:
         logging.error(e)
         return False
+    return True
+
+def ping_server(aris_name, state):
+    """Upload a notification file to AWS
+
+    :param aris_name: Name of the aris file uploaded
+    :return: True if file was uploaded, else False
+    """
+
+    file_name = 'tmp/notification.txt'
+    os.makedirs('tmp', exist_ok=True)
+    with open(file_name, 'w') as f:
+        output = f"time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
+        output += f"filename: {aris_name}\n"
+        output += f"app version: {state['version']}\n"
+        output += f"hardware: {torch.cuda.get_device_name() if torch.cuda.is_available() else 'CPU'}\n"
+        if 'CPU_CORES' in os.environ:
+            output += f"location: HuggingFace\n"
+            output += f"nbr cpu cores: {os.environ['CPU_CORES']}\n"
+            output += f"memory: {os.environ['MEMORY']}\n"
+        else:
+            output += f"location: local install"
+        f.write(output)
+ 
+    # If S3 object_name was not specified, use file_name
+
+    if (not 'AAK_ID' in os.environ) or (not 'ASAK' in os.environ):
+        print('AWS keys not specified. Cancelling sync')
+        return False
+
+    # Upload the file
+    s3_client = boto3.client(
+        's3',
+        aws_access_key_id=os.environ['AAK_ID'],
+        aws_secret_access_key=os.environ['ASAK']
+    )
+    try:
+        response = s3_client.upload_file(file_name, "fishcounting", "webapp_uploads/notifications/" + str(int(datetime.now().timestamp())) + ".txt")
+    except ClientError as e:
+        logging.error(e)
+        return False
     return True

dataloader.py → backend/dataloader.py

@@ -5,20 +5,19 @@ import cv2
 import numpy as np
 import json
 from threading import Lock
-import struct
 from contextlib import contextmanager
 import torch
 from torch.utils.data import Dataset
-import torchvision.transforms as T
 from PIL import Image
+from datetime import datetime
 
 # assumes yolov5 on sys.path
 from lib.yolov5.utils.general import xyxy2xywh
 from lib.yolov5.utils.augmentations import letterbox
 from lib.yolov5.utils.dataloaders import create_dataloader as create_yolo_dataloader
 
-from pyDIDSON import pyDIDSON
-from aris import ImageData
+from backend.pyDIDSON import pyDIDSON
+from backend.aris import BEAM_WIDTH_DIR
 
 # use this flag to test the difference between direct ARIS dataloading and
 # using the jpeg compressed version. very slow. not much difference observed.
@@ -29,21 +28,24 @@ TEST_JPG_COMPRESSION = False
 # Factory(ish) methods for DataLoader creation. Easy entry points to this module.
 # # # # # #
 
-def create_dataloader_aris(aris_filepath, beam_width_dir, annotations_file, batch_size=32, stride=64, pad=0.5, img_size=896, rank=-1, world_size=1, workers=0, 
-                      disable_output=False, cache_bg_frames=False):
+def create_dataloader_aris(aris_filepath, beam_width_dir=BEAM_WIDTH_DIR, annotations_file=None, batch_size=32, stride=64, pad=0.5, img_size=896, rank=-1, world_size=1, workers=0, 
+                      disable_output=False, cache_bg_frames=False, num_frames_bg_subtract=1000):
     """
     Get a PyTorch Dataset and DataLoader for ARIS files with (optional) associated fisheye-formatted labels.
     """
+
+    print('dataset', datetime.now())
     # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
     # this is a no-op for a single-gpu machine
     with torch_distributed_zero_first(rank):
         dataset = YOLOARISBatchedDataset(aris_filepath, beam_width_dir, annotations_file, stride, pad, img_size, 
-                                         disable_output=disable_output, cache_bg_frames=cache_bg_frames)
+                                         disable_output=disable_output, cache_bg_frames=cache_bg_frames, num_frames_bg_subtract=num_frames_bg_subtract)
 
     
     batch_size = min(batch_size, len(dataset))
     nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers])  # number of workers
     
+    print('dataloader', datetime.now())
     if not disable_output:
         print("dataset size", len(dataset))
         print("dataset shape", dataset.shape)
@@ -55,6 +57,7 @@ def create_dataloader_aris(aris_filepath, beam_width_dir, annotations_file, batc
                                                         num_workers=nw,
                                                         pin_memory=True,
                                                         collate_fn=collate_fn)
+    print('done', datetime.now())
     return dataloader, dataset
 
 def create_dataloader_frames(frames_path, batch_size=32, model_stride_max=32, 
@@ -348,8 +351,8 @@ class YOLOARISBatchedDataset(ARISBatchedDataset):
     """An ARIS Dataset that works with YOLOv5 inference."""
     
     def __init__(self, aris_filepath, beam_width_dir, annotations_file, stride=64, pad=0.5, img_size=896, batch_size=32, 
-                 disable_output=False, cache_bg_frames=False):
-        super().__init__(aris_filepath, beam_width_dir, annotations_file, batch_size, disable_output=disable_output, cache_bg_frames=cache_bg_frames)
+                 disable_output=False, cache_bg_frames=False, num_frames_bg_subtract=1000):
+        super().__init__(aris_filepath, beam_width_dir, annotations_file, batch_size, disable_output=disable_output, cache_bg_frames=cache_bg_frames, num_frames_bg_subtract=num_frames_bg_subtract)
         
         # compute shapes for letterbox
         aspect_ratio = self.ydim / self.xdim

inference.py → backend/inference.py

@@ -5,21 +5,21 @@ from tqdm import tqdm
 from functools import partial
 import numpy as np
 import json
-import time
 from unittest.mock import patch
 import math
+import torch
+import torchvision
 
 # assumes yolov5 on sys.path
 from lib.yolov5.models.experimental import attempt_load
 from lib.yolov5.utils.torch_utils import select_device
 from lib.yolov5.utils.general import clip_boxes, scale_boxes, xywh2xyxy
 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.bytetrack import Associate
+from lib.fish_eye.tracker_bytetrack import Associate
+
+from backend.InferenceConfig import InferenceConfig, TrackerType
 
 
 ### Configuration options
@@ -51,56 +51,59 @@ def norm(bbox, w, h):
     bb[3] /= h
     return bb
 
-def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None, config=InferenceConfig()):
+def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None, config=InferenceConfig(), verbose=True):
     
     # Set up model
     model, device = setup_model(config.weights)
 
     # Detect boxes in frames
-    inference, image_shapes, width, height = do_detection(dataloader, model, device, gp=gp)
+    inference, image_shapes, width, height = do_detection(dataloader, model, device, gp=gp, verbose=verbose)
 
+    result = do_full_tracking(inference, image_shapes, image_meter_width, image_meter_height, width, height, config=InferenceConfig(), gp=None, verbose=verbose)
+    
+    return result
+
+def do_full_tracking(inference, image_shapes, image_meter_width, image_meter_height, width, height, config=InferenceConfig(), gp=None, verbose=True):
     if config.associative_tracker == TrackerType.BYTETRACK:
 
         # Find low confidence detections
-        low_outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.byte_low_conf, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp)
-        low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, gp=gp)
+        low_outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.byte_low_conf, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp, verbose=verbose)
+        low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, gp=gp, verbose=verbose)
 
         # Find high confidence detections
-        high_outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.byte_high_conf, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp)
-        high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, gp=gp)
+        high_outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.byte_high_conf, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp, verbose=verbose)
+        high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, gp=gp, verbose=verbose)
 
         # Perform associative tracking (ByteTrack)
         results = do_associative_tracking(
             low_preds, high_preds, image_meter_width, image_meter_height, 
             reverse=False, min_length=config.min_length, min_travel=config.min_travel, 
             max_age=config.max_age, min_hits=config.min_hits, 
-            gp=gp)
+            gp=gp, verbose=verbose)
     else: 
 
         # Find confident detections
-        outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp)
+        outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp, verbose=verbose)
 
         if config.associative_tracker == TrackerType.CONF_BOOST:
 
             # Boost confidence based on found confident detections
-            do_confidence_boost(inference, outputs, boost_power=config.boost_power, boost_decay=config.boost_decay, gp=gp)
+            do_confidence_boost(inference, outputs, boost_power=config.boost_power, boost_decay=config.boost_decay, gp=gp, verbose=verbose)
 
             # Find confident detections from boosted list
-            outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp)
+            outputs = do_suppression(inference, image_meter_width, width, conf_thres=config.conf_thresh, iou_thres=config.nms_iou, max_length=config.max_length, gp=gp, verbose=verbose)
 
         # Format confident detections
-        all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, gp=gp)
+        all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, gp=gp, verbose=verbose)
 
         # Perform SORT tracking
         results = do_tracking(
             all_preds, image_meter_width, image_meter_height, 
             min_length=config.min_length, min_travel=config.min_travel, 
             max_age=config.max_age, min_hits=config.min_hits, 
-            gp=gp)
+            gp=gp, verbose=verbose)
  
     return results
-    
-
 
 
 

main.py → backend/predict.py

@@ -1,11 +1,13 @@
+import project_path
+
 import os
 import torch
 from zipfile import ZipFile
 
-from aris import create_manual_marking, BEAM_WIDTH_DIR, add_metadata_to_result, prep_for_mm
-from dataloader import create_dataloader_aris
-from inference import do_full_inference, json_dump_round_float
-from visualizer import generate_video_batches
+from backend.aris import create_manual_marking, BEAM_WIDTH_DIR, add_metadata_to_result, prep_for_mm
+from backend.dataloader import create_dataloader_aris
+from backend.inference import do_full_inference, json_dump_round_float
+from backend.visualizer import generate_video_batches
 
 def predict_task(filepath, config, output_formats=[], gradio_progress=None):
     """
@@ -32,6 +34,7 @@ def predict_task(filepath, config, output_formats=[], gradio_progress=None):
 
     # Create dataloader
     if (gradio_progress): gradio_progress(0, "Initializing Dataloader...")
+    dataloader, dataset = create_dataloader_aris(filepath, BEAM_WIDTH_DIR, None, num_frames_bg_subtract=0)
     dataloader, dataset = create_dataloader_aris(filepath, BEAM_WIDTH_DIR, None)
 
     # Extract aris/didson info. Didson does not yet have pixel-meter info
@@ -55,11 +58,11 @@ def predict_task(filepath, config, output_formats=[], gradio_progress=None):
     json_dump_round_float(results, results_filepath)
     
     # Create Manual Marking file
-    if "Manual Marking" in output_formats and dataset.didson.info['version'][3] == 5: 
+    if "Generate Manual Marking" in output_formats and dataset.didson.info['version'][3] == 5: 
         create_manual_marking(results, out_path=marking_filepath)
     
     # Create Annotated Video
-    if "Annotated Video" in output_formats:
+    if "Generate Annotated Video" in output_formats:
         generate_video_batches(dataset.didson, results, frame_rate, video_filepath, 
                    image_meter_width=image_meter_width, image_meter_height=image_meter_height, gp=gradio_progress)
     

pyDIDSON.py → backend/pyDIDSON.py

@@ -17,7 +17,7 @@ import struct
 from types import SimpleNamespace
 
 import lib.fish_eye.pyARIS as pyARIS
-from pyDIDSON_format import *
+from backend.pyDIDSON_format import *
 
 
 class pyDIDSON:

uploader.py → backend/uploader.py

@@ -1,4 +1,3 @@
-import project_path
 import os;
 from datetime import datetime;
 
@@ -30,12 +29,6 @@ def save_data_to_dir(bytes, filename, dirname):
         return False, None, None
     
     return True, filepath, dirname
-    
-
-def allowed_file(filename):
-    """Only allow an ARIS file to be uploaded."""
-    return '.' in filename and \
-           filename.rsplit('.', 1)[1].lower() in ['aris', 'ddf']
 
 def create_data_dir(identifier = None):
     """Create a (probably) unique directory for a task."""

visualizer.py → backend/visualizer.py

@@ -1,6 +1,5 @@
 import project_path
 
-import json
 import cv2
 import numpy as np
 from tqdm import tqdm

{gradio_scripts → frontend}/annotation_handler.py

@@ -1,13 +1,75 @@
-import json
 import cv2
 import base64
+import gradio as gr
+import json
+import numpy as np
 
 VIDEO_HEIGHT = 700
 
 
+# annotation_btn.clock - switches to annotation tab and starts load_annotation
+def prepare_annotation(state, result, result_index):
+
+    state['annotation_index'] = result_index
+    state['frame_index'] = 0
+
+    # output for [annotation_progress, master_tabs]
+    if result["aris_input"][result_index]:
+        return [
+            gr.update(value="<p id='annotation_info' style='display:none'>[]</p><!--" + str(np.random.rand()) + "-->", visible=True),
+            gr.update(selected=2)
+        ]
+    return [gr.update(), gr.update()]
+
+# annotation_progress.change - loads annotation frames in batches - called after prepare_annotation
+def load_annotation(state, result, progress_bar):
+
+    # Get result index
+    result_index = state['annotation_index']
+
+    set_progress = lambda pct, msg: progress_bar(pct, desc=msg)
+
+    if state['frame_index'] == 0:
+        if set_progress: set_progress(0, "Loading Frames")
+
+    # Check that frames remain to be loaded
+    if state['frame_index'] < len(result['json_result'][result_index]['frames']):
+
+        # load frames and annotation
+        annotation_info, state['frame_index'] = init_frames(result["aris_input"][result_index], result['json_result'][result_index], state['frame_index'], gp=set_progress)
+    
+        # save as html element
+        annotation_content = "<p id='annotation_info' style='display:none'>" + json.dumps(annotation_info) + "</p>"
+
+        # output for [annotation_editor, annotation_progress]
+        return [gr.update(), gr.update(value=annotation_content)]
+
+    # If complete, start annotation editor    
+
+    annotation_html = ""
+
+    # Header
+    annotation_html += "<div id='annotation_header'>"
+    annotation_html += "     <h1 id='annotation_frame_nbr'>Frame 0/100</h1>"
+    annotation_html += "     <p id='annotation_edited'>(edited)</p>"
+    annotation_html += "</div>"
+
+    # Annotation Body
+    annotation_html += "<div style='display:flex'>"
+    annotation_html += "     <canvas id='canvas' style='width:50%' onmousedown='mouse_down(event)' onmousemove='mouse_move(event)' onmouseup='mouse_up()' onmouseleave='mouse_up()'></canvas>"
+    annotation_html += "     <div id='annotation_display' style='width:50%'></div>"
+    annotation_html += "</div>"
+    
+    # Dummy objects
+    annotation_html += "<img id='annotation_img' onload='draw()' style='display:none'></img>"
+    annotation_html += "<!--" + str(np.random.rand()) + "-->"
+
+    # output for [annotation_editor, annotation_progress]
+    return [gr.update(value=annotation_html, visible=True), gr.update(visible=False)]
+
+# called by load_annotation - read frames from dataloader and formats tracks
 def init_frames(dataset, preds, index, gp=None):
     """Load frames for annotation editing
-
     
         Returns:
             list({
@@ -19,6 +81,7 @@ def init_frames(dataset, preds, index, gp=None):
                 )
             })
     """
+
     images = dataset.didson.load_frames(start_frame=0, end_frame=1)
 
     # assumes all frames the same size
@@ -39,6 +102,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=index+i, end_frame=index+i+1)[0]
             image = cv2.resize(cv2.cvtColor(img_raw, cv2.COLOR_GRAY2BGR), (w, h))
             retval, buffer = cv2.imencode('.jpg', image)
@@ -65,3 +129,43 @@ def init_frames(dataset, preds, index, gp=None):
     
     return annotations, end_index
 
+# javascript code that retrieves the data from load_annotation and saves it to the javascript window
+js_store_frame_info = """
+    () => {
+        info_string = document.getElementById("annotation_info").innerHTML;
+        info = JSON.parse(info_string);
+        console.log(info)
+        if (info.length == 0) {
+            window.annotation_info = [];
+            return false;
+        } 
+        window.annotation_info = window.annotation_info.concat(info)
+        console.log(window.annotation_info)
+        return true;
+    }
+"""
+
+annotation_css = """
+#annotation_frame_nbr {
+    left: calc(50% - 100px);
+    position: absolute;
+    width: 200px;
+    text-align: center;
+    font-size: x-large;
+}
+#annotation_header {
+    height: 40px;
+}
+#annotation_frame_nbr {
+    left: calc(50% - 100px);
+    position: absolute;
+    width: 200px;
+    text-align: center;
+    font-size: x-large;
+}
+#annotation_edited {
+    right: 0px;
+    position: absolute;
+    margin-top: 5px;
+}
+"""

frontend/aris_crop.py

@@ -0,0 +1,45 @@
+import lib.fish_eye.pyARIS as pyARIS
+import struct
+
+def crop_clip(aris_path, num_frames, verbose=False):
+    """
+    Crop an aris file based on the first *num_frames* frames. Save the new aris files in the tmp folder
+    """
+
+    # load aris file and extract frame size
+    ARIS_data, frame = pyARIS.DataImport(aris_path)
+    FrameSize = ARIS_data.NumRawBeams*ARIS_data.SamplesPerChannel
+    if verbose: print("True Old", ARIS_data.FrameCount, ARIS_data.StartFrame, ARIS_data.EndFrame)
+
+    # get byte index of the cutoff point
+    frameoffset = (1024+(num_frames*(1024+(FrameSize))))
+
+    # read aris the bytes for the head and frames we want and cast to bytearray
+    data = open(ARIS_data.filename, 'rb')
+    cropped = data.read(frameoffset)
+    array = bytearray(cropped)
+
+    # get old values for important metadata
+    old_frame_count     = struct.unpack("I", array[4:8])[0]
+    old_start_frame     = struct.unpack("I", array[352:356])[0]
+    old_end_frame       = struct.unpack("I", array[356:360])[0]
+
+    if verbose: print("old", old_frame_count, old_start_frame, old_end_frame)
+
+    # set new values
+    array[4:8]        = bytearray(struct.pack("I", num_frames))
+    array[352:356]    = bytearray(struct.pack("I", old_start_frame))
+    array[356:360]    = bytearray(struct.pack("I", old_start_frame + num_frames))
+    if verbose: print("new", array[4:8], array[352:356], array[356:360])
+    
+    # cast to bytes
+    cropped = bytes(array)
+
+    # save new aris file
+    with open("tmp/cropped_aris.aris", 'wb') as f:
+        f.write(cropped)
+
+    # load file to check that the frame count, start frame and end frame makes sense
+    if verbose: 
+        ARIS_data_2, frame = pyARIS.DataImport("tmp/cropped_aris.aris")
+        print("check", ARIS_data_2.FrameCount, ARIS_data_2.StartFrame, ARIS_data_2.EndFrame)

{gradio_scripts → frontend}/pdf_handler.py

@@ -1,26 +1,23 @@
 import datetime
 import numpy as np
 from matplotlib.backends.backend_pdf import PdfPages
-from matplotlib import collections as mc
 import matplotlib.pyplot as plt
 import math
-from aris import BEAM_WIDTH_DIR
 import cv2
-
-from dataloader import create_dataloader_aris
-
+import os
 
 STANDARD_FIG_SIZE = (16, 9)
-OUT_PDF_FILE_NAME = 'multipage_pdf.pdf'
+OUT_PDF_FILE_NAME = 'tmp/fisheye_pdf.pdf'
+os.makedirs('tmp', exist_ok=True)
 
 
-def make_pdf(i, state, result, table_headers):
+def make_pdf(i, state, result, dataset, table_headers):
 
     fish_info = result["fish_info"][i]
     fish_table = result["fish_table"][i]
     json_result = result['json_result'][i]
+    dataset = result['datasets'][i]
     metadata = json_result['metadata']
-    aris_input = result["aris_input"][i]
 
     with PdfPages(OUT_PDF_FILE_NAME) as pdf:
         plt.rcParams['text.usetex'] = False
@@ -31,11 +28,6 @@ def make_pdf(i, state, result, table_headers):
 
         generate_fish_list(pdf, table_headers, fish_table)
 
-
-        dataset = None
-        if (aris_input is not None):
-            dataloader, dataset = create_dataloader_aris(aris_input, BEAM_WIDTH_DIR, None)
-
         for i, fish in enumerate(json_result['fish']):
             calculate_fish_paths(json_result, dataset, i)
         
@@ -173,6 +165,7 @@ def calculate_fish_paths(result, dataset, id):
     img = None
     if (dataset is not None):
         
+
         images = dataset.didson.load_frames(start_frame=start_frame, end_frame=start_frame+1)
         img = images[0]
 
@@ -358,10 +351,6 @@ def draw_fish_tracks(pdf, result, dataset, id):
 
 
 
-
-
-
-
     if (dataset is not None):
         indices = [start_frame, int(2/3*start_frame + end_frame/3), int(1/3*start_frame + 2/3*end_frame), end_frame]
         fig, axs = plt.subplots(2, len(indices), sharex=False, sharey=False, figsize=STANDARD_FIG_SIZE)
@@ -377,15 +366,6 @@ def draw_fish_tracks(pdf, result, dataset, id):
                         box = ann['bbox']
                         frame_index = fi
                         break
-
-            batch_i = math.floor(frame_index/32)
-            fi = frame_index - batch_i*32
-            batch = dataset[batch_i]
-            (rgb_img, _, shapes) = batch[fi]
-            rgb_img = rgb_img.permute(1, 2, 0)
-            print(type(batch))
-            print(type(rgb_img))
-            print(rgb_img.shape)
             
             print("box", i, box)
             if box is not None:
@@ -400,18 +380,6 @@ def draw_fish_tracks(pdf, result, dataset, id):
                 axs[0, i].imshow(cropped_img, extent=(cx-s, cx+s, cy-s, cy+s), cmap=plt.colormaps['Greys_r'])
                 axs[0, i].plot([x1, x1, x2, x2, x1], [y1, y2, y2, y1, y1], color="red")
                 axs[0, i].set_title('Frame ' + str(frame_index))
-
-                h, w, _ = rgb_img.shape
-                print(w, h)
-                x1, x2, y1, y2 = int(box[0]*w), int(box[2]*w), int(box[1]*h), int(box[3]*h)
-                cx, cy = int((x2 + x1)/2), int((y2 + y1)/2)
-                s = min(int(max(x2 - x1, y2 - y1)*5/2), cx, cy, w-cx, h-cy)
-                print(x1, x2, y1, y2)
-                print(cx, cy, s)
-                cropped_img = rgb_img[cy-s:cy+s, cx-s:cx+s, :]
-                axs[1, i].imshow(cropped_img, extent=(cx-s, cx+s, cy-s, cy+s), cmap=plt.colormaps['Greys_r'])
-                axs[1, i].plot([x1, x1, x2, x2, x1], [y1, y2, y2, y1, y1], color="red")
-                axs[1, i].set_title('Frame ' + str(frame_index))
         
         pdf.savefig(fig)
         plt.close(fig)

{gradio_scripts → frontend}/result_ui.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import numpy as np
-from gradio_scripts.pdf_handler import make_pdf
 import os
 
+from frontend.pdf_handler import make_pdf
+
 js_update_tab_labels = """
     async () => {
         let el_list = document.getElementById("tab_labeler").getElementsByClassName("svelte-1kcgrqr")
@@ -39,9 +40,9 @@ def update_result(i, state, result, inference_handler):
     # Check if inference is done
     not_done = state['index'] < state['total']
 
-    annotation_avaliable = not (result["aris_input"][i] == None)
+    annotation_avaliable = state['enable_annotation_editor'] and (result["aris_input"][i] is not None)
 
-    if 'PDF' in state['outputs']:
+    if 'Generate PDF' in state['outputs']:
         print("making pdf")
         make_pdf(state['index']-1, state, result, table_headers)
         print("done pdf")
@@ -65,9 +66,14 @@ def update_result(i, state, result, inference_handler):
     }
 
 
+# Auto_download
+def auto_download_zip(state):
+    if 'Automatically download result' in state['outputs']:
+        return gr.update(value=str(np.random.rand()))
+    else:
+        return gr.update()
 
-
-def Result_Gradio(prepare_annotation, components):
+def Result_Gradio(prepare_annotation, components, state):
     global tabs, tab_parent, zip_out
 
     # Dummy element to call inference events, this also displays the inference progress
@@ -86,9 +92,23 @@ def Result_Gradio(prepare_annotation, components):
     visual_components = []
 
     # Zip file output
-    zip_out = gr.File(label="ZIP Output", interactive=False)
+    zip_out = gr.File(label="ZIP Output", elem_id="zip_out", interactive=False)
     visual_components.append(zip_out)
-
+    components['zip_out'] = zip_out
+
+    autodownloader = gr.Text(value="LOADING", visible=False)
+
+    zip_out.change(lambda: auto_download_zip(state), None, autodownloader)
+    autodownloader.change(lambda x: x, autodownloader, None, _js="""
+        () => {
+            zip_out = document.getElementById("zip_out")
+            downloads = zip_out?.getElementsByClassName("download")
+            if (downloads?.length > 0) {
+                downloads[downloads.length-1].children[0].click()       
+            }
+        }
+        """
+    )
 
     # Create result tabs
     tabs = []
@@ -128,4 +148,48 @@ def Result_Gradio(prepare_annotation, components):
 
     components['result_tabs'] = tab_parent
 
-    return visual_components
+    return visual_components
+
+
+def create_metadata_table(result, table_headers, info_headers):
+    if 'metadata' in result:
+        metadata = result['metadata']
+    else:
+        metadata = { 'FISH': [] }
+    
+    # Calculate detection dropout
+    for fish in metadata['FISH']:
+        count = 0
+        for frame in result['frames'][fish['START_FRAME']:fish['END_FRAME']+1]:
+            for ann in frame['fish']:
+                if ann['fish_id'] == fish['TOTAL']:
+                    count += 1
+        fish['DETECTION_DROPOUT'] = 1 - count / (fish['END_FRAME'] + 1 - fish['START_FRAME'])
+
+    # Create fish table
+    table = []
+    for fish in metadata["FISH"]:
+        row = []
+        for header in table_headers:
+            row.append(fish[header])
+        table.append(row)
+    
+    if len(metadata["FISH"]) == 0:
+        row = []
+        for header in table_headers:
+            row.append("-")
+        table.append(row)
+
+    # Create info table
+    info = []
+    for field in info_headers:
+        field_name = "**" + field + "**"
+        if field in metadata:
+            info.append([field_name, str(metadata[field])])
+        else:
+            info.append([field_name, ""])
+    if 'hyperparameters' in metadata:
+        for param_name in metadata['hyperparameters']:
+            info.append(['**' + param_name + '**', str(metadata['hyperparameters'][param_name])])
+    
+    return table, info

frontend/state_handler.py

@@ -0,0 +1,22 @@
+
+def reset_state(result, state):
+
+    # Reset Result
+    result["json_result"] = []
+    result["aris_input"] = []
+    result["datasets"] = []
+    result["path_video"] = []
+    result["path_zip"] = []
+    result["path_json"] = []
+    result["path_marking"] = []
+    result["fish_table"] = []
+    result["fish_info"] = []
+
+    # Reset State
+    state['files'] = []
+    state['index'] = 0
+    state['total'] = 0
+
+
+
+

{gradio_scripts → frontend}/upload_ui.py

@@ -1,13 +1,13 @@
 import gradio as gr
-from gradio_scripts.file_reader import File
-from InferenceConfig import InferenceConfig, TrackerType
-
+from frontend.custom_file_reader import File
+from backend.InferenceConfig import InferenceConfig, TrackerType
+import os
 
 models = {
     'master': 'models/v5m_896_300best.pt',
-    'elwha': 'models/YsEE20.pt',
-    'elwha+kenai_val': 'models/YsEKvE20.pt',
-    'elwha+kenai_train': 'models/YsEKtE20.pt',
+#    'elwha': 'models/YsEE20.pt',
+#    'elwha+kenai_val': 'models/YsEKvE20.pt',
+    'elwha': 'models/YsEKtE20.pt',
 }
 
 def Upload_Gradio(gradio_components):
@@ -38,28 +38,33 @@ def Upload_Gradio(gradio_components):
                 tracker = gr.Dropdown(["None", "Confidence Boost", "ByteTrack"], value=default_tracker, label="Associative Tracking")
                 hyperparams.append(tracker)
                 with gr.Row(visible=default_tracker=="Confidence Boost") as track_row:
-                    hyperparams.append(gr.Slider(0, 5, value=default_settings.boost_power, label="Boost Power", info=""))
-                    hyperparams.append(gr.Slider(0, 1, value=default_settings.boost_decay, label="Boost Decay", info=""))
+                    hyperparams.append(gr.Slider(0, 5, value=default_settings.boost_power, label="Boost Power", info="Scalar multiplier for the boost amount"))
+                    hyperparams.append(gr.Slider(0, 1, value=default_settings.boost_decay, label="Boost Decay", info="Exponential decay parameter for boost based on frame time difference"))
                     tracker.change(lambda x:  gr.update(visible=(x=="Confidence Boost")), tracker, track_row)
                 with gr.Row(visible=default_tracker=="ByteTrack") as track_row:
-                    hyperparams.append(gr.Slider(0, 1, value=default_settings.byte_low_conf, label="Low Conf Threshold", info=""))
-                    hyperparams.append(gr.Slider(0, 1, value=default_settings.byte_high_conf, label="High Conf Threshold", info=""))
+                    hyperparams.append(gr.Slider(0, 1, value=default_settings.byte_low_conf, label="Low Conf Threshold", info="Confidence threshold for the low detection group"))
+                    hyperparams.append(gr.Slider(0, 1, value=default_settings.byte_high_conf, label="High Conf Threshold", info="Confidence threshold for the high detection group"))
                     tracker.change(lambda x:  gr.update(visible=(x=="ByteTrack")), tracker, track_row)
 
                 gr.Markdown("Other")
                 with gr.Row():
                     hyperparams.append(gr.Slider(0, 3, value=default_settings.min_length, label="Min Length", info="Minimum length of fish (meters) in order for it to count"))
-                    hyperparams.append(gr.Slider(0, 3, value=default_settings.max_length, label="Max Length", info="Maximum length of fish (meters) in order for it to count"))
+                    hyperparams.append(gr.Slider(0, 3, value=default_settings.max_length, label="Max Length", info="Maximum length of fish (meters) in order for it to count. (disable at 0)"))
                     hyperparams.append(gr.Slider(0, 10, value=default_settings.min_travel, label="Min Travel", info="Minimum travel distance of track (meters) in order for it to count"))
 
                 gradio_components['hyperparams'] = hyperparams
 
             with gr.Row():
-                hyperparams.append(gr.CheckboxGroup(["Annotated Video", "Manual Marking", "PDF"], label="Output formats", interactive=True, value=["Annotated Video", "Manual Marking"]))
+                hyperparams.append(gr.CheckboxGroup([("Generate Annotated Video"), ("Generate Manual Marking"), ("Generate PDF"), ("Automatically download result")], label="Output settings", interactive=True, value=["Generate Annotated Video"]))
 
             #Input field for aris submission
             gradio_components['input'] = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")
 
+            example_name = "static/example.aris"
+            gradio_components['examples'] = gr.Examples(examples=[[example_name]], inputs=gradio_components['input'])
+
+            gradio_components['inference_btn'] = gr.Button("Run")
+
         # Tab - uploading old result files to review
         with gr.Tab("Open Result"):
             gr.HTML("""

gradio_scripts/state_handler.py

@@ -1,451 +0,0 @@
-from aris import create_metadata_table
-import json
-
-example_metadata = {
-    "FILE_NAME":  "static/example_metadata/fisheye",
-    "FRAME_RATE":  6.548702716827393,
-    "UPSTREAM_FISH":  0,
-    "DOWNSTREAM_FISH":  0,
-    "NONDIRECTIONAL_FISH":  14,
-    "TOTAL_FISH":  14,
-    "TOTAL_FRAMES":  644,
-    "EXPECTED_FRAMES":  -1,
-    "TOTAL_TIME":  "0:01:38",
-    "EXPECTED_TIME":  "0:00:00",
-    "UPSTREAM_MOTION":  "Right To Left",
-    "COUNT_FILE_NAME":  "N/A",
-    "EDITOR_ID":  "N/A",
-    "INTENSITY":  "0.0 dB",
-    "THRESHOLD":  "0.0 dB",
-    "WINDOW_START":  1,
-    "WINDOW_END":  17,
-    "WATER_TEMP":  "13 degC",
-    "FISH":  [
-        {
-            "FILE":  1,
-            "TOTAL":  1,
-            "FRAME_NUM":  12,
-            "DIR":  " N/A",
-            "R":  13.403139282569885,
-            "THETA":  0.1706,
-            "L":  63.739999999999995,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:54:40",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  2,
-            "FRAME_NUM":  35,
-            "DIR":  " N/A",
-            "R":  13.206211097755432,
-            "THETA":  -9.1195,
-            "L":  73.33,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:54:44",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  3,
-            "FRAME_NUM":  122,
-            "DIR":  " N/A",
-            "R":  13.219339643409729,
-            "THETA":  -9.3961,
-            "L":  84.77,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:54:58",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  4,
-            "FRAME_NUM":  123,
-            "DIR":  "N/A",
-            "R":  12.996154367286682,
-            "THETA":  10.7991,
-            "L":  59.919999999999995,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:54:58",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  5,
-            "FRAME_NUM":  130,
-            "DIR":  " N/A",
-            "R":  12.484141086769105,
-            "THETA":  -8.2654,
-            "L":  70.89999999999999,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:54:59",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  6,
-            "FRAME_NUM":  218,
-            "DIR":  " N/A",
-            "R":  13.232468189064026,
-            "THETA":  -9.3961,
-            "L":  77.25999999999999,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:55:12",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  7,
-            "FRAME_NUM":  278,
-            "DIR":  " N/A",
-            "R":  13.967666745704651,
-            "THETA":  -12.8758,
-            "L":  37.51,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:55:22",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  8,
-            "FRAME_NUM":  302,
-            "DIR":  " N/A",
-            "R":  13.25872528037262,
-            "THETA":  -9.1195,
-            "L":  79.5,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:55:25",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  9,
-            "FRAME_NUM":  331,
-            "DIR":  " N/A",
-            "R":  13.25872528037262,
-            "THETA":  -9.1195,
-            "L":  80.67,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:55:30",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  10,
-            "FRAME_NUM":  450,
-            "DIR":  " N/A",
-            "R":  13.324368008644104,
-            "THETA":  -8.5535,
-            "L":  83.1,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:55:48",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  11,
-            "FRAME_NUM":  495,
-            "DIR":  " N/A",
-            "R":  13.481910556495666,
-            "THETA":  -9.1195,
-            "L":  86.39,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:55:55",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  12,
-            "FRAME_NUM":  526,
-            "DIR":  " N/A",
-            "R":  13.04866854990387,
-            "THETA":  10.5397,
-            "L":  55.37,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:56:00",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  13,
-            "FRAME_NUM":  538,
-            "DIR":  " N/A",
-            "R":  13.416267828224182,
-            "THETA":  -9.668,
-            "L":  82.38,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:56:02",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }, {
-            "FILE":  1,
-            "TOTAL":  14,
-            "FRAME_NUM":  624,
-            "DIR":  " N/A",
-            "R":  13.29811091733551,
-            "THETA":  -8.8385,
-            "L":  77.44,
-            "DR":  -1,
-            "LDR":  -1,
-            "ASPECT":  -1,
-            "TIME":  "11:56:16",
-            "DATE":  "2018-07-09",
-            "LATITUDE":  "N 00 d 0.00000 m",
-            "LONGITUDE":  "E 000 d 0.00000 m",
-            "PAN":  None,
-            "TILT":  None,
-            "ROLL":  0,
-            "SPECIES":  "Unknown",
-            "MOTION":  "Running <-->",
-            "Q":  -1,
-            "N":  -1,
-            "COMMENT":  ""
-        }
-    ],
-    "DATE":  "2018-07-09",
-    "START":  "11:54:39",
-    "END":  "11:56:18"
-}
-
-def load_example_result(result, table_headers, info_headers):
-    fish_table, fish_info = create_metadata_table(example_metadata, table_headers, info_headers)
-    result['path_zip'] = ["static/example/input_file_results.zip"]
-    result['path_video'] = ["static/example/input_file_results.mp4"]
-    result['path_json'] = ["static/example/input_file_results.json"]
-    result['path_marking'] = ["static/example/input_file_marking.txt"]
-    result['fish_table'] = [fish_table]
-    result['fish_info'] = [fish_info]
-
-
-def reset_state(result, state):
-
-    # Reset Result
-    result["json_result"] = []
-    result["aris_input"] = []
-    result["path_video"] = []
-    result["path_zip"] = []
-    result["path_json"] = []
-    result["path_marking"] = []
-    result["fish_table"] = []
-    result["fish_info"] = []
-
-    # Reset State
-    state['files'] = []
-    state['index'] = 0
-    state['total'] = 0
-
-
-def convert_json_to_vatic(json_path, vatic_path="static/example/input_file_vatic.xml"):
-    
-    xml = '<?xml version="1.0" encoding="utf-8"?>\n';
-    xml += '<annotation>\n';
-    xml += '  <folder>not available</folder>\n';
-    xml += '  <filename>not available</filename>\n';
-    xml += '  <source>\n';
-    xml += '    <type>video</type>\n';
-    xml += '    <sourceImage>vatic frames</sourceImage>\n';
-    xml += '    <sourceAnnotation>vatic</sourceAnnotation>\n';
-    xml += '  </source>\n';
-
-    with open(json_path, 'r') as f:
-        annotation = json.loads(f.read())
-
-    frames = annotation['frames']
-    nbr_frames = len(frames)
-
-    fishes = {}
-    for frame in annotation['frames']:
-        frame_nbr = str(frame['frame_num'])
-        for fish in frame['fish']: 
-            track_id = fish['fish_id']
-            if (not track_id in fishes): fishes[track_id] = {'id': track_id, 'frames': []}
-
-            fishes[track_id]['frames'].append({
-                'frame': frame_nbr,
-                'x_min': str(round(fish['bbox'][0]*522)),
-                'y_min': str(round(fish['bbox'][1]*700)),
-                'x_max': str(round(fish['bbox'][2]*522)),
-                'y_max': str(round(fish['bbox'][3]*700)),
-                'visible': str(fish['visible']),
-                'truth': "1"
-            })
-
-    for fish_id in fishes:
-        fish = fishes[fish_id]
-
-        xml += '  <object>\n';
-        xml += '    <name>fish</name>\n';
-        xml += '    <moving>true</moving>\n';
-        xml += '    <action/>\n';
-        xml += '    <verified>0</verified>\n';
-        xml += '    <id>' + str(fish['id']) + '</id>\n';
-        xml += '    <createdFrame>0</createdFrame>\n';
-        xml += '    <startFrame>0</startFrame>\n';
-        xml += '    <endFrame>' + str(nbr_frames - 1 ) + '</endFrame>\n';
-    
-        for frame in fish['frames']:
-            xml += '    ';
-            xml += '<polygon>';
-            xml += '<t>' + str(frame['frame']) + '</t>';
-            xml += '<pt><x>' + frame['x_min'] + '</x><y>' + frame['y_min'] + '</y><l>' + frame['truth'] + '</l></pt>';
-            xml += '<pt><x>' + frame['x_min'] + '</x><y>' + frame['y_max'] + '</y><l>' + frame['truth'] + '</l></pt>';
-            xml += '<pt><x>' + frame['x_max'] + '</x><y>' + frame['y_max'] + '</y><l>' + frame['truth'] + '</l></pt>';
-            xml += '<pt><x>' + frame['x_max'] + '</x><y>' + frame['y_min'] + '</y><l>' + frame['truth'] + '</l></pt>';
-            xml += '</polygon>\n';
-        
-        xml += '  </object>\n';
-
-    xml += '</annotation>\n';
-
-    if vatic_path:
-        with open(vatic_path, 'w') as f:
-            f.write(xml)
-    return xml
-
-
-convert_json_to_vatic("static/example/input_file_results.json")
-
-

lib/fish_eye/.gitignore

@@ -1,104 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-.hypothesis/
-.pytest_cache/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# pyenv
-.python-version
-
-# celery beat schedule file
-celerybeat-schedule
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/

lib/fish_eye/tracker.py

@@ -5,8 +5,8 @@ import json
 import numpy as np
 
 from fish_length import Fish_Length
-from lib.fish_eye.sort import Sort
-from lib.fish_eye.bytetrack import Associate
+from lib.fish_eye.tracker_sort import Sort
+from lib.fish_eye.tracker_bytetrack import Associate
 import lib
 
 class Tracker:
@@ -30,13 +30,13 @@ class Tracker:
             # Match confidence with correct track
             conf = 0
             min_score = 1000000
-            if type(self.algorithm) == lib.fish_eye.sort.Sort:
+            if type(self.algorithm) == lib.fish_eye.tracker_sort.Sort:
                 for det in dets:
                     score = sum(abs(det[0:4] - track[0:4]))
                     if (score < min_score):
                         min_score = score
                         conf = det[4]
-            elif type(self.algorithm) == lib.fish_eye.bytetrack.Associate:
+            elif type(self.algorithm) == lib.fish_eye.tracker_bytetrack.Associate:
                 for det in dets[0]:
                     score = sum(abs(det[0:4] - track[0:4]))
                     if (score < min_score):

lib/fish_eye/{sort.py → tracker_sort.py}

@@ -214,7 +214,7 @@ class Sort(object):
         trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
         for t in reversed(to_del):
             self.trackers.pop(t)
-        matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets, trks, self.iou_threshold)
+        matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets,trks, self.iou_threshold)
 
         # update matched trackers with assigned detections
         for m in matched:

scripts/{infer_aris.py → aris_to_tracks.py}

@@ -1,13 +1,20 @@
-import project_path
 import argparse
-from predict import predict_task
-from datetime import datetime
+
+import project_subpath
+from backend.InferenceConfig import InferenceConfig
+from backend.predict import predict_task
 
 
 def main(args):
-    predict_task(args.aris, weights=args.weights)
+    predict_task(args.aris, config=InferenceConfig())
 
 def argument_parser():
+    """
+    Run full inference on an aris file
+    Args:
+        aris (str): Path to ARIS file. Required.
+        weights (str): Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument("--aris", required=True, help="Path to ARIS file. Required.")
     parser.add_argument("--weights", default='../models/v5m_896_300best.pt', help="Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt")

scripts/{track_detection.py → detection_to_tracks.py}

@@ -1,59 +1,38 @@
-import project_path
-from lib.yolov5.utils.torch_utils import select_device
-from lib.yolov5.utils.general import clip_boxes, scale_boxes
 import argparse
-from datetime import datetime
 import torch
 import os
-from dataloader import create_dataloader_frames_only
-from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking, do_associative_tracking
-from visualizer import generate_video_batches
 import json
 from tqdm import tqdm
-import numpy as np
 
+import project_subpath
 
-def main(args, config={}, verbose=True):
+from backend.InferenceConfig import InferenceConfig
+from backend.inference import do_full_tracking
+
+
+def main(args, config=InferenceConfig(), verbose=True):
     """
-    Main processing task to be run in gradio
-        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
-        - Writes json output to dirname(filepath)/{filename}_results.json
-        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
-        - Writes video output to dirname(filepath)/{filename}_results.mp4
-        - Zips all results to dirname(filepath)/{filename}_results.zip
+    Convert raw detections to tracks and saves the tracking json result
     Args:
-        filepath (str): path to aris file
-        
-    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+        detections (str): path to raw detections directory. Required
+        output (str): where tracking result will be stored. Required
+        metadata (str): path to metadata directory. Required
+        tracker (str): arbitrary name of tracker folder that you want to save trajectories to
     """
 
-    # setup config
-    if "conf_threshold" not in config: config['conf_threshold'] = 0.3#0.001
-    if "nms_iou" not in config: config['nms_iou'] = 0.3#0.6
-    if "min_length" not in config: config['min_length'] = 0.3
-    if "min_travel" not in config: config['min_travel'] = 0
-    if "max_age" not in config: config['max_age'] = 20
-    if "iou_threshold" not in config: config['iou_threshold'] = 0.01
-    if "min_hits" not in config: config['min_hits'] = 11
-    if "associativity" not in config: config['associativity'] = None
-
-    print(config)
+    print("running detections_to_tracks.py with:", config.to_dict())
 
-    
-    locations = [
-        "kenai-rightbank"
-    ]
-    for loc in locations:
+    loc = args.location
 
-        in_loc_dir = os.path.join(args.detections, loc)
-        out_loc_dir = os.path.join(args.output, loc, args.tracker, "data")
-        os.makedirs(out_loc_dir, exist_ok=True)
-        metadata_path = os.path.join(args.metadata, loc + ".json")
-        print(in_loc_dir)
-        print(out_loc_dir)
-        print(metadata_path)
+    in_loc_dir = os.path.join(args.detections, loc)
+    out_loc_dir = os.path.join(args.output, loc, args.tracker, "data")
+    os.makedirs(out_loc_dir, exist_ok=True)
+    metadata_path = os.path.join(args.metadata, loc + ".json")
+    print(in_loc_dir)
+    print(out_loc_dir)
+    print(metadata_path)
 
-        track_location(in_loc_dir, out_loc_dir, metadata_path, config, verbose)
+    track_location(in_loc_dir, out_loc_dir, metadata_path, config, verbose)
 
 
                 
@@ -100,28 +79,14 @@ def track(in_loc_dir, out_loc_dir, metadata_path, seq, config, verbose):
                 image_meter_height = sequence['y_meter_start'] - sequence['y_meter_stop']
 
 
-    if config['associativity'] == "bytetrack":
-
-        low_outputs = do_suppression(inference, conf_thres=config['low_conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-        low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, verbose=verbose)
-
-        high_outputs = do_suppression(inference, conf_thres=config['high_conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-        high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, verbose=verbose)
-
-        results = do_associative_tracking(low_preds, high_preds, image_meter_width, image_meter_height, reverse=False, min_length=config['min_length'], min_travel=config['min_travel'], max_age=config['max_age'], iou_thres=config['iou_threshold'], min_hits=config['min_hits'], verbose=verbose)
-    else:
-        outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-
-        if config['associativity'] == "boost":
-
-            do_confidence_boost(inference, outputs, boost_power=config['boost_power'], boost_decay=config['boost_decay'], verbose=verbose)
-
-            outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-
-        all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, verbose=verbose)
+    # assume all images in the sequence have the same shape
+    real_width = image_shapes[0][0][0][1]
+    real_height = image_shapes[0][0][0][0]
 
-        results = do_tracking(all_preds, image_meter_width, image_meter_height, min_length=config['min_length'], min_travel=config['min_travel'], max_age=config['max_age'], iou_thres=config['iou_threshold'], min_hits=config['min_hits'], verbose=verbose)
+    # perform tracking
+    results = do_full_tracking(inference, image_shapes, image_meter_width, image_meter_height, width, height, config=config, gp=None, verbose=verbose)
 
+    # write tracking result
     mot_rows = []
     for frame in results['frames']:
         for fish in frame['fish']:
@@ -152,9 +117,10 @@ def track(in_loc_dir, out_loc_dir, metadata_path, seq, config, verbose):
 def argument_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--detections", required=True, help="Path to frame directory. Required.")
+    parser.add_argument("--location", required=True, help="Name of location dir. Required.")
     parser.add_argument("--output", required=True, help="Path to output directory. Required.")
-    parser.add_argument("--metadata", required=True, help="Path to output directory. Required.")
-    parser.add_argument("--tracker", default='tracker', help="Path to output directory. Required.")
+    parser.add_argument("--metadata", required=True, help="Path to metadata directory. Required.")
+    parser.add_argument("--tracker", default='tracker', help="Tracker name.")
     return parser
 
 if __name__ == "__main__":

scripts/detection_to_tracks_eval.py

@@ -0,0 +1,86 @@
+import argparse
+import sys, os
+import numpy as np
+
+import project_subpath
+from backend.InferenceConfig import InferenceConfig
+
+from detection_to_tracks import main as track
+
+current_dir = os.path.dirname(os.path.realpath(__file__))
+pardir = os.path.dirname(current_dir)
+sys.path.append(os.path.join(pardir, "../caltech-fish-counting/"))
+from evaluate import evaluate
+
+class Object(object):
+    pass
+
+def main(args):
+    """
+    Perform tracking on a directory of raw detections, saves the tracks, and runs the 'evaluate' script from the 'caltech-fish-counting' repo
+    Args:
+        detection_dir (str): path to raw detection directory
+        weights (str): path to weights
+        conf_threshold (float): confidence cutoff for detection filtering
+        nms_iou (float): non-maximum suppression IOU threshold
+        min_length (float): minimum length of fish in meters in order to count
+        max_length (float): maximum length of fish in meters in order to count. Disable with 0
+        min_travel (float): minimum travel distance in meters of track in order to count
+        max_age (int): aximum time between detections before a fish is forgotten by the tracker
+        min_hits (int): minimum length of track in frames in order to count
+        associativity (str): string representation of tracking method with corresponding hyperparameters separated by ':'
+        verbose (bool): disable or enable logging
+    """
+
+    infer_args = Object()
+    infer_args.detections = args.detection_dir
+    infer_args.metadata = "../frames/metadata"
+    infer_args.output = "../frames/result_testing"
+    infer_args.tracker = 'tracker'
+    infer_args.location = 'kenai-val'
+
+    config = InferenceConfig(
+        conf_thresh=float(args.conf_threshold),
+        nms_iou=float(args.nms_iou),
+        min_hits=int(args.min_hits),
+        max_age=int(args.max_age),
+        min_length=float(args.min_length),
+        max_length=float(args.max_length),
+        min_travel=float(args.min_travel),
+    )
+
+    config.enable_tracker_from_string(args.associativity)
+
+    print("verbose", args.verbose)
+
+    track(infer_args, config=config, verbose=args.verbose)
+
+    result = evaluate(infer_args.output, "../frames/MOT", "../frames/metadata", infer_args.tracker, True, location=infer_args.location)
+    metrics = result['MotChallenge2DBox']['tracker']['COMBINED_SEQ']['pedestrian']
+    print('HOTA:', np.mean(metrics['HOTA']['HOTA'])*100)
+    print('MOTA:', metrics['CLEAR']['MOTA']*100)
+    print('IDF1:', metrics['Identity']['IDF1']*100)
+    print('nMAE:', metrics['nMAE']['nMAE']*100)
+    print('misscounts:', str(metrics['nMAE']['nMAE_numer']) + "/" + str(metrics['nMAE']['nMAE_denom']))
+    return result
+
+
+def argument_parser():
+    default = InferenceConfig()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--detection_dir", default="../frames/detection_storage", help="Path to raw detection directory")
+    parser.add_argument("--weights", default=default.weights,               help="Path to weights")
+    parser.add_argument("--conf_threshold", default=default.conf_thresh,    help="Confidence cutoff for detection filtering")
+    parser.add_argument("--nms_iou", default=default.nms_iou,               help="Non-maximum Suppression IOU threshold")
+    parser.add_argument("--min_length", default=default.min_length,         help="Minimum length of fish in meters in order to count")
+    parser.add_argument("--max_length", default=default.max_length,         help="Maximum length of fish in meters in order to count. Disable with 0")
+    parser.add_argument("--min_travel", default=default.min_travel,         help="Minimum travel distance in meters of track in order to count.")
+    parser.add_argument("--max_age", default=default.max_age,               help="Maximum time between detections before a fish is forgotten by the tracker")
+    parser.add_argument("--min_hits", default=default.min_hits,             help="Minimum length of track in frames in order to count")
+    parser.add_argument("--associativity", default='',                      help="String representation of tracking method with corresponding hyperparameters separated by ':'")
+    parser.add_argument("--verbose", action='store_true',                   help="Disable or enable logging")
+    return parser
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)

scripts/{full_detect_frames.py → frames_to_MOT.py}

@@ -1,66 +1,62 @@
-import project_path
-from lib.yolov5.utils.general import clip_boxes, scale_boxes
 import argparse
-from datetime import datetime
 import torch
 import os
-from dataloader import create_dataloader_frames_only
-from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking
-from visualizer import generate_video_batches
 import json
 from tqdm import tqdm
-import numpy as np
 
+import project_subpath
+from backend.dataloader import create_dataloader_frames_only
+from backend.inference import setup_model, do_detection, do_suppression
+from backend.InferenceConfig import InferenceConfig
+
+from lib.yolov5.utils.general import clip_boxes, scale_boxes
 
-def main(args, config={}, verbose=True):
+
+def main(args, config=InferenceConfig(), verbose=False):
     """
-    Main processing task to be run in gradio
-        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
-        - Writes json output to dirname(filepath)/{filename}_results.json
-        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
-        - Writes video output to dirname(filepath)/{filename}_results.mp4
-        - Zips all results to dirname(filepath)/{filename}_results.zip
+    Construct and save MOT format detections from yolov5 based on a frame directory
     Args:
-        filepath (str): path to aris file
-        
-    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+        frames (str): path to image directory
+        output (str): where MOT detections will be stored
+        weights (str): path to model weights
     """
+    
     print("In task...")
     print("Cuda available in task?", torch.cuda.is_available())
 
-    # setup config
-    if "conf_threshold" not in config: config['conf_threshold'] = 0.001
-    if "nms_iou" not in config: config['nms_iou'] = 0.6
-    if "min_length" not in config: config['min_length'] = 0.3
-    if "max_age" not in config: config['max_age'] = 20
-    if "iou_threshold" not in config: config['iou_threshold'] = 0.01
-    if "min_hits" not in config: config['min_hits'] = 11
-
-    print(config)
+    print("Config:", config.to_dict())
 
     model, device = setup_model(args.weights)
-    
-    locations = [
-        "kenai-val"
-    ]
-    for loc in locations:
 
-        in_loc_dir = os.path.join(args.frames, loc)
-        out_loc_dir = os.path.join(args.output, loc)
-        print(in_loc_dir)
-        print(out_loc_dir)
+    in_loc_dir = os.path.join(args.frames, args.location)
+    out_loc_dir = os.path.join(args.output, args.location)
+    metadata_path = os.path.join(args.metadata, args.location + ".json")
+    print(in_loc_dir)
+    print(out_loc_dir)
+    print(metadata_path)
 
-        detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose)
+    detect_location(in_loc_dir, out_loc_dir, metadata_path, config, model, device, verbose)
 
 
                 
-def detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose):
+def detect_location(in_loc_dir, out_loc_dir, metadata_path, config, model, device, verbose):
 
     seq_list = os.listdir(in_loc_dir)
 
     with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
         for seq in seq_list:
 
+
+            image_meter = (-1, -1)
+            with open(metadata_path, 'r') as f:
+                json_object = json.loads(f.read())
+                for sequence in json_object:
+                    if sequence['clip_name'] == seq:
+                        image_meter = (
+                            sequence['x_meter_stop'] - sequence['x_meter_start'],
+                            sequence['y_meter_stop'] - sequence['y_meter_start']
+                        )
+
             pbar.update(1)
             if (seq.startswith(".")): continue
             pbar.set_description("Processing " + seq)
@@ -69,12 +65,12 @@ def detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose):
             out_seq_dir = os.path.join(out_loc_dir, seq)
             os.makedirs(out_seq_dir, exist_ok=True)
 
-            detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose)
+            detect_seq(in_seq_dir, out_seq_dir, image_meter, config, model, device, verbose)
 
-def detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose):
+def detect_seq(in_seq_dir, out_seq_dir, image_meter, config, model, device, verbose):
 
     ann_list = []
-    frame_list = detect(in_seq_dir, config, model, device, verbose)
+    frame_list = detect(in_seq_dir, image_meter, config, model, device, verbose)
     for frame in frame_list:
         if frame is not None:
             for ann in frame:
@@ -89,7 +85,7 @@ def detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose):
     with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
         f.write(result)
 
-def detect(in_dir, config, model, device, verbose):
+def detect(in_dir, image_meter, config, model, device, verbose):
     
     #progress_log = lambda p, m: 0
 
@@ -98,8 +94,7 @@ def detect(in_dir, config, model, device, verbose):
 
     inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
 
-
-    outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
+    outputs = do_suppression(inference, image_meter_width=image_meter[0], image_pixel_width=image_meter[1], conf_thres=config.conf_thresh, iou_thres=config.nms_iou, verbose=verbose)
 
     file_names = dataloader.files
     frame_list = []
@@ -130,6 +125,8 @@ def detect(in_dir, config, model, device, verbose):
 def argument_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--frames", required=True, help="Path to frame directory. Required.")
+    parser.add_argument("--metadata", required=True, help="Path to frame directory. Required.")
+    parser.add_argument("--location", required=True, help="Name of location dir. Required.")
     parser.add_argument("--output", required=True, help="Path to output directory. Required.")
     parser.add_argument("--weights", default='models/v5m_896_300best.pt', help="Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt")
     return parser

scripts/{detect_frames.py → frames_to_detections.py}

@@ -1,46 +1,35 @@
-import project_path
-from lib.yolov5.utils.general import clip_boxes, scale_boxes
 import argparse
-from datetime import datetime
 import torch
 import os
-from dataloader import create_dataloader_frames_only
-from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking
-from visualizer import generate_video_batches
 import json
 from tqdm import tqdm
-import numpy as np
 
+import project_subpath
+from backend.dataloader import create_dataloader_frames_only
+from backend.inference import setup_model, do_detection
 
-def main(args, verbose=True):
+
+def main(args, verbose=False):
     """
-    Main processing task to be run in gradio
-        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
-        - Writes json output to dirname(filepath)/{filename}_results.json
-        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
-        - Writes video output to dirname(filepath)/{filename}_results.mp4
-        - Zips all results to dirname(filepath)/{filename}_results.zip
+    Construct and save raw detections from yolov5 based on a frame directory
     Args:
-        filepath (str): path to aris file
-        
-    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+        frames (str): path to image directory
+        output (str): where detections will be stored
+        weights (str): path to model weights
     """
+    
+
     print("In task...")
     print("Cuda available in task?", torch.cuda.is_available())
 
     model, device = setup_model(args.weights)
-    
-    locations = [
-        "kenai-rightbank"
-    ]
-    for loc in locations:
 
-        in_loc_dir = os.path.join(args.frames, loc)
-        out_loc_dir = os.path.join(args.output, loc)
-        print(in_loc_dir)
-        print(out_loc_dir)
+    in_loc_dir = os.path.join(args.frames, args.location)
+    out_loc_dir = os.path.join(args.output, args.location)
+    print(in_loc_dir)
+    print(out_loc_dir)
 
-        detect_location(in_loc_dir, out_loc_dir, model, device, verbose)
+    detect_location(in_loc_dir, out_loc_dir, model, device, verbose)
 
 
                 
@@ -82,6 +71,7 @@ def detect(in_seq_dir, out_seq_dir, model, device, verbose):
 def argument_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--frames", default="../frames/images", help="Path to frame directory. Required.")
+    parser.add_argument("--location", default="kenai-val", help="Name of location dir. Required.")
     parser.add_argument("--output", default="../frames/detections/detection_storage/", help="Path to output directory. Required.")
     parser.add_argument("--weights", default='models/v5m_896_300best.pt', help="Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt")
     return parser

scripts/{infer_frames.py → frames_to_tracks.py}

@@ -1,70 +1,58 @@
-import project_path
+
 import argparse
-from datetime import datetime
 import torch
 import os
-from dataloader import create_dataloader_frames_only
-from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking
-from visualizer import generate_video_batches
 import json
 from tqdm import tqdm
 
+import project_subpath
+from backend.InferenceConfig import InferenceConfig
+from backend.dataloader import create_dataloader_frames_only
+from backend.inference import do_full_tracking, setup_model, do_detection
+
 
-def main(args, config={}, verbose=True):
+def main(args, config=InferenceConfig(), verbose=True):
     """
-    Main processing task to be run in gradio
-        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
-        - Writes json output to dirname(filepath)/{filename}_results.json
-        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
-        - Writes video output to dirname(filepath)/{filename}_results.mp4
-        - Zips all results to dirname(filepath)/{filename}_results.zip
+    Perform inference on a directory of frames and saves the tracking json result
     Args:
-        filepath (str): path to aris file
-        
-    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+        frames (str): Path to frame directory. Required.
+        metadata (str): Path to metadata directory. Required.
+        output (str): Path to output directory. Required.
+        weights (str): Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt
     """
+    
     print("In task...")
     print("Cuda available in task?", torch.cuda.is_available())
 
-    # setup config
-    if "conf_threshold" not in config: config['conf_threshold'] = 0.3
-    if "nms_iou" not in config: config['nms_iou'] = 0.3
-    if "min_length" not in config: config['min_length'] = 0.3
-    if "max_age" not in config: config['max_age'] = 20
-    if "iou_threshold" not in config: config['iou_threshold'] = 0.01
-    if "min_hits" not in config: config['min_hits'] = 11
-
-    print(config)
+    print("Config:", config.to_dict())
 
     dirname = args.frames
-    
-    locations = ["kenai-val"]
-    for loc in locations:
-
-        in_loc_dir = os.path.join(dirname, loc)
-        out_dir = os.path.join(args.output, loc, "tracker", "data")
-        metadata_path = os.path.join(args.metadata, loc + ".json")
-        os.makedirs(out_dir, exist_ok=True)
-        print(in_loc_dir)
-        print(out_dir)
-        print(metadata_path)
-
-        # run detection + tracking
-        model, device = setup_model(args.weights)
-
-        seq_list = os.listdir(in_loc_dir)
-        idx = 1
-        with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
-            for seq in seq_list:
-                pbar.update(1)
-                pbar.set_description("Processing " + seq)
-                if verbose:
-                    print(" ")
-                    print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
-                    print(" ")
-                idx += 1
-                in_seq_dir = os.path.join(in_loc_dir, seq)
-                infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path, verbose)
+    loc = args.location
+
+    in_loc_dir = os.path.join(dirname, loc)
+    out_dir = os.path.join(args.output, loc, "tracker", "data")
+    metadata_path = os.path.join(args.metadata, loc + ".json")
+    os.makedirs(out_dir, exist_ok=True)
+    print(in_loc_dir)
+    print(out_dir)
+    print(metadata_path)
+
+    # run detection + tracking
+    model, device = setup_model(args.weights)
+
+    seq_list = os.listdir(in_loc_dir)
+    idx = 1
+    with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
+        for seq in seq_list:
+            pbar.update(1)
+            pbar.set_description("Processing " + seq)
+            if verbose:
+                print(" ")
+                print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
+                print(" ")
+            idx += 1
+            in_seq_dir = os.path.join(in_loc_dir, seq)
+            infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path, verbose)
 
 def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path, verbose):
     
@@ -94,17 +82,11 @@ def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path, v
             f.write("ERROR")
         return
 
+    real_width = image_shapes[0][0][0][1]
+    real_height = image_shapes[0][0][0][0]
 
-    outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-
-    do_confidence_boost(inference, outputs, verbose=verbose)
-
-    new_outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-
-    all_preds, real_width, real_height = format_predictions(image_shapes, new_outputs, width, height)
-
-    results = do_tracking(all_preds, image_meter_width, image_meter_height, min_length=config['min_length'], max_age=config['max_age'], iou_thres=config['iou_threshold'], min_hits=config['min_hits'], verbose=verbose)
-
+    results = do_full_tracking(inference, image_shapes, image_meter_width, image_meter_height, width, height, config=config, gp=None, verbose=verbose)
+    
     mot_rows = []
     for frame in results['frames']:
         for fish in frame['fish']:
@@ -137,6 +119,7 @@ def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path, v
 def argument_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--frames", required=True, help="Path to frame directory. Required.")
+    parser.add_argument("--location", required=True, help="Name of location dir. Required.")
     parser.add_argument("--metadata", required=True, help="Path to metadata directory. Required.")
     parser.add_argument("--output", required=True, help="Path to output directory. Required.")
     parser.add_argument("--weights", default='models/v5m_896_300best.pt', help="Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt")

scripts/frames_to_tracks_eval.py

@@ -0,0 +1,78 @@
+import argparse
+import sys, os
+import numpy as np
+
+import project_subpath
+from frames_to_tracks import main as infer
+from backend.InferenceConfig import InferenceConfig
+
+current_dir = os.path.dirname(os.path.realpath(__file__))
+pardir = os.path.dirname(current_dir)
+sys.path.append(os.path.join(pardir, "../caltech-fish-counting/"))
+from evaluate import evaluate
+
+class Object(object):
+    pass
+
+def main(args):
+    """
+    Perform inference on a directory of frames, saves the tracks, and runs the 'evaluate' script from the 'caltech-fish-counting' repo
+    Args:
+        weights (str): path to weights
+        conf_threshold (float): confidence cutoff for detection filtering
+        nms_iou (float): non-maximum suppression IOU threshold
+        min_length (float): minimum length of fish in meters in order to count
+        max_length (float): maximum length of fish in meters in order to count. Disable with 0
+        min_travel (float): minimum travel distance in meters of track in order to count
+        max_age (int): aximum time between detections before a fish is forgotten by the tracker
+        min_hits (int): minimum length of track in frames in order to count
+        associativity (str): string representation of tracking method with corresponding hyperparameters separated by ':'
+    """
+
+    infer_args = Object()
+    infer_args.metadata = "../frames/metadata"
+    infer_args.frames = "../frames/images"
+    infer_args.location = "kenai-val"
+    infer_args.output = "../frames/result"
+    infer_args.weights = "models/v5m_896_300best.pt"
+
+    config = InferenceConfig(
+        conf_thresh=float(args.conf_threshold),
+        nms_iou=float(args.nms_iou),
+        min_hits=int(args.min_hits),
+        max_age=int(args.max_age),
+        min_length=float(args.min_length),
+        max_length=float(args.max_length),
+        min_travel=float(args.min_travel),
+    )
+
+    config.enable_tracker_from_string(args.associativity)
+
+    infer(infer_args, config=config, verbose=False)
+
+    result = evaluate("../frames/result", "../frames/MOT", "../frames/metadata", "tracker", True, location=infer_args.location)
+    metrics = result['MotChallenge2DBox']['tracker']['COMBINED_SEQ']['pedestrian']
+    print('HOTA:', np.mean(metrics['HOTA']['HOTA'])*100)
+    print('MOTA:', metrics['CLEAR']['MOTA']*100)
+    print('IDF1:', metrics['Identity']['IDF1']*100)
+    print('nMAE:', metrics['nMAE']['nMAE']*100)
+    print('misscounts:', str(metrics['nMAE']['nMAE_numer']) + "/" + str(metrics['nMAE']['nMAE_denom']))
+    return result
+
+def argument_parser():
+    default = InferenceConfig()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--weights", default=default.weights,               help="Path to weights")
+    parser.add_argument("--conf_threshold", default=default.conf_thresh,    help="Confidence cutoff for detection filtering")
+    parser.add_argument("--nms_iou", default=default.nms_iou,               help="Non-maximum Suppression IOU threshold")
+    parser.add_argument("--min_length", default=default.min_length,         help="Minimum length of fish in meters in order to count")
+    parser.add_argument("--max_length", default=default.max_length,         help="Maximum length of fish in meters in order to count. Disable with 0")
+    parser.add_argument("--min_travel", default=default.min_travel,         help="Minimum travel distance in meters of track in order to count.")
+    parser.add_argument("--max_age", default=default.max_age,               help="Maximum time between detections before a fish is forgotten by the tracker")
+    parser.add_argument("--min_hits", default=default.min_hits,             help="Minimum length of track in frames in order to count")
+    parser.add_argument("--associativity", default='',                      help="String representation of tracking method with corresponding hyperparameters separated by ':'")
+    return parser
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)

scripts/infer_eval.py

@@ -1,47 +0,0 @@
-import project_path
-import argparse
-from infer_frames import main as infer
-import sys
-sys.path.append('..')
-sys.path.append('../caltech-fish-counting')
-
-from evaluate import evaluate
-
-class Object(object):
-    pass
-
-def main(args):
-
-    infer_args = Object()
-    infer_args.metadata = "../frames/metadata"
-    infer_args.frames = "../frames/images"
-    infer_args.output = "../frames/result"
-    infer_args.weights = "models/v5m_896_300best.pt"
-
-    config = {
-        'conf_threshold': float(args.conf_threshold),
-        'nms_iou': float(args.nms_iou),
-        'min_length': float(args.min_length),
-        'max_age': int(args.max_age),
-        'iou_threshold': float(args.iou_threshold),
-        'min_hits': int(args.min_hits)
-    }
-
-    infer(infer_args, config=config, verbose=False)
-
-    evaluate("../frames/result", "../frames/MOT", "../frames/metadata", "tracker", False)
-
-
-def argument_parser():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--conf_threshold", default=0.3, help="Config object. Required.")
-    parser.add_argument("--nms_iou", default=0.3, help="Config object. Required.")
-    parser.add_argument("--min_length", default=0.3, help="Config object. Required.")
-    parser.add_argument("--max_age", default=20, help="Config object. Required.")
-    parser.add_argument("--iou_threshold", default=0.01, help="Config object. Required.")
-    parser.add_argument("--min_hits", default=11, help="Config object. Required.")
-    return parser
-
-if __name__ == "__main__":
-    args = argument_parser().parse_args()
-    main(args)

scripts/{project_path.py → project_subpath.py}

@@ -3,6 +3,8 @@ import os
 
 current_dir = os.path.dirname(os.path.realpath(__file__))
 pardir = os.path.dirname(current_dir)
-for d in [pardir, current_dir, os.path.join(pardir, "lib/fish_eye/"), os.path.join(pardir, "lib/"), os.path.join(pardir, "lib/yolov5/")]:
+for d in [pardir, current_dir]:
     if d not in sys.path:
-        sys.path.append(d)
+        sys.path.append(d)
+
+import project_path

scripts/track_eval.py

@@ -1,78 +0,0 @@
-import project_path
-import argparse
-from track_detection import main as track
-import sys
-import numpy as np
-sys.path.append('..')
-sys.path.append('../caltech-fish-counting')
-
-from evaluate import evaluate
-
-class Object(object):
-    pass
-
-def main(args):
-
-    infer_args = Object()
-    infer_args.detections = args.detection_dir
-    infer_args.metadata = "../frames/metadata"
-    infer_args.output = "../frames/result_testing"
-    infer_args.tracker = 'tracker'
-
-    config = {
-        'conf_threshold': float(args.conf_threshold),
-        'nms_iou': float(args.nms_iou),
-        'min_length': float(args.min_length),
-        'min_travel': float(args.min_travel),
-        'max_age': int(args.max_age),
-        'iou_threshold': float(args.iou_threshold),
-        'min_hits': int(args.min_hits),
-        'associativity': None
-    }
-
-    if args.associativity != "":
-        if (args.associativity.startswith("boost")):
-            config['associativity'] = "boost"
-            conf = args.associativity.split(":")
-            if len(conf) > 1: config['boost_power'] = float(conf[1])
-            if len(conf) > 2: config['boost_decay'] = float(conf[2])
-        elif (args.associativity.startswith("bytetrack")):
-            config['associativity'] = "bytetrack"
-            conf = args.associativity.split(":")
-            if len(conf) > 1: config['low_conf_threshold'] = float(conf[1])
-            if len(conf) > 2: config['high_conf_threshold'] = float(conf[2])
-        else:
-            print("INVALID ASSOCIATIVITY TYPE:", args.associativity)
-            return
-
-    print("verbose", args.verbose)
-
-    track(infer_args, config=config, verbose=args.verbose)
-
-    result = evaluate(infer_args.output, "../frames/MOT", "../frames/metadata", infer_args.tracker, True)
-    metrics = result['MotChallenge2DBox']['tracker']['COMBINED_SEQ']['pedestrian']
-    print('HOTA:', np.mean(metrics['HOTA']['HOTA'])*100)
-    print('MOTA:', metrics['CLEAR']['MOTA']*100)
-    print('IDF1:', metrics['Identity']['IDF1']*100)
-    print('nMAE:', metrics['nMAE']['nMAE']*100)
-    print('misscounts:', str(metrics['nMAE']['nMAE_numer']) + "/" + str(metrics['nMAE']['nMAE_denom']))
-    return result
-
-
-def argument_parser():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--detection_dir", default="../frames/detection_storage")
-    parser.add_argument("--conf_threshold", default=0.3, help="Config object. Required.")
-    parser.add_argument("--nms_iou", default=0.3, help="Config object. Required.")
-    parser.add_argument("--min_length", default=0.3, help="Config object. Required.")
-    parser.add_argument("--min_travel", default=0, help="Config object. Required.")
-    parser.add_argument("--max_age", default=20, help="Config object. Required.")
-    parser.add_argument("--iou_threshold", default=0.01, help="Config object. Required.")
-    parser.add_argument("--min_hits", default=11, help="Config object. Required.")
-    parser.add_argument("--associativity", default='', help="Config object. Required.")
-    parser.add_argument("--verbose", action='store_true', help="Config object. Required.")
-    return parser
-
-if __name__ == "__main__":
-    args = argument_parser().parse_args()
-    main(args)