Detection only script

app.py

@@ -1,21 +1,16 @@
 import gradio as gr
 from uploader import save_data_to_dir, create_data_dir, save_data
 from main import predict_task
-from state_handler import load_example_result, reset_state
-from visualizer import is_fourcc_available
-from file_reader import File
+from gradio_scripts.state_handler import reset_state
 import numpy as np
-from aws_handler import upload_file
+from gradio_scripts.aws_handler import upload_file
 from aris import create_metadata_table
-import cv2
-import base64
-from bbox import draggable_js
-from annotation_handler import load_frames
+from gradio_scripts.annotation_handler import load_frames
 import json
 from zipfile import ZipFile
 import os
-from gradio_components.upload_ui import Upload_Gradio
-from gradio_components.result_ui import Result_Gradio, update_result, table_headers, info_headers, js_update_tab_labels
+from gradio_scripts.upload_ui import Upload_Gradio
+from gradio_scripts.result_ui import Result_Gradio, update_result, table_headers, info_headers, js_update_tab_labels
 
 
 #Initialize State & Result
@@ -46,6 +41,12 @@ def on_aris_input(file_list):
 # Called when a result zip file is uploaded for result review
 def on_result_upload(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)
 
     component_updates = { 
@@ -92,7 +93,6 @@ def on_result_upload(zip_list, aris_list):
         result["fish_table"].append(fish_table)
         result["fish_info"].append(fish_info)
 
-        print(result['aris_input'])
         update = update_result(i, state, result, inference_handler)
 
         for key in update.keys():
@@ -190,10 +190,19 @@ def open_annotation(index):
     if result["aris_input"][index]:
         frame_info = load_frames(result["aris_input"][index], result['json_result'][index])
 
-        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>"
+        # 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 += "<p id='annotation_info' style='display:none'>" + json.dumps(frame_info) + "</p>"
         annotation_html += "<img id='annotation_img' onload='draw()' style='display:none'></img>"
         annotation_html += "<!--" + str(np.random.rand()) + "-->"
@@ -216,7 +225,29 @@ with demo:
                     display: none !important;
                 }
                 .selected.svelte-kqij2n {
-                    background: linear-gradient(180deg, #66eecb47, white);
+                    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;
                 }
             </style>
             <style id="tab_style"></style>
@@ -257,7 +288,7 @@ with demo:
                 annotation_progress.change(open_annotation, annotation_progress, [annotation_editor, annotation_progress], _js="() => window.annotation_index")
                 
                 # Event listener for running javascript defined in 'annotation_editor.js'
-                with open('annotation_editor.js', 'r') as f:
+                with open('gradio_scripts/annotation_editor.js', 'r') as f:
                     annotation_editor.change(lambda x: gr.update(), None, annotation_editor, _js=f.read())
 
     # Disclaimer at the bottom of page

aris.py

@@ -15,8 +15,8 @@ from copy import deepcopy
 from multiprocessing import Pool
 import math
 
-import pyARIS
-from tracker import Tracker
+import lib.fish_eye.pyARIS as pyARIS
+from lib.fish_eye.tracker import Tracker
 
 
 BEAM_WIDTH_DIR = 'lib/fish_eye/beam_widths/'

bbox.py

@@ -1,309 +0,0 @@
-draggable_js = """
-() => {
-    window.canvas = document.getElementById('canvas');
-    window.ctx = canvas.getContext('2d');
-    window.rects = [];
-    window.mouseX;
-    window.mouseY;
-    window.closeEnough = 10;
-    window.keys = {};
-    window.hover = false;
-
-    window.TL = 0;
-    window.TR = 1;
-    window.BL = 2;
-    window.BR = 3;
-
-
-    window.init = () => {
-        window.frames = JSON.parse(document.getElementById("annotation_info").innerHTML);
-        window.frame_index = 0;
-
-        document.addEventListener('keydown', keydown);
-        document.addEventListener('keyup', keyup);
-
-        show_frame();
-    }
-    window.prev_frame = () => {
-        window.frame_index = Math.max(window.frame_index - 1, 0);
-        show_frame();
-    }
-    window.next_frame = () => {
-        window.frame_index = Math.min(window.frame_index + 1, window.frames.length - 1);
-        show_frame();
-    }
-    window.show_frame = () => {
-        const frame_info = window.frames[window.frame_index];
-        const annotations = frame_info['annotations'];
-        const frame = frame_info['frame'];
-
-
-        window.annotations = [];
-        for (const annotation of annotations) {
-            const bbox = annotation['bbox']
-            window.annotations.push({
-                rect: {
-                    startX: bbox.left,
-                    startY: bbox.top,
-                    w: bbox.right - bbox.left,
-                    h: bbox.bottom - bbox.top
-                },
-                id: annotation.id,
-                conf: annotation.conf
-            })
-        }
-        console.log(window.annotations)
-        window.dragging = false;
-
-        document.getElementById("annotation_img").src = "data:image/png;base64," + frame;
-    }
-
-    window.draw = () => {
-        ctx = window.ctx;
-        canvas = window.canvas;
-        canvas.width = document.getElementById("annotation_img").width;
-        canvas.height = document.getElementById("annotation_img").height;
-        canvas.style = ""
-        annotations = window.annotations;
-    
-        ctx.clearRect(0, 0, canvas.width, canvas.height);
-        
-        ctx.drawImage(document.getElementById("annotation_img"), 0, 0);
-
-        for (const annotation of annotations) {
-            //ctx.globalAlpha = annotation.conf
-
-            const rect = annotation.rect;
-            ctx.strokeStyle = color_from_id(annotation.id);
-            ctx.strokeRect(rect.startX, rect.startY, rect.w, rect.h);
-
-            ctx.font = "15px Arial";
-            ctx.fillStyle = color_from_id(annotation.id);
-            ctx.textAlign = "right";
-            ctx.fillText(annotation.id, rect.startX + rect.w, rect.startY - 3);
-        }
-
-        if (hover && !dragging) {
-            annotation = hover.annotation;
-            rect = annotation.rect;
-            handles = [
-                [rect.startX, rect.startY], 
-                [rect.startX + rect.w, rect.startY], 
-                [rect.startX, rect.startY + rect.h], 
-                [rect.startX + rect.w, rect.startY + rect.h]
-            ];
-
-            handle = handles[hover.corner];
-            ctx.fillStyle = color_from_id(annotation.id);
-            ctx.beginPath();
-            s = 6;
-            ctx.rect(handle[0]-s/2, handle[1]-s/2, s, s);
-            ctx.fill();
-        }
-
-        prettify_annotation();
-    }
-    color_from_id = (id) => {
-        return 'hsl(' + Math.floor((id*id)*57 % 360) + ', 100%, 50%)'
-    }
-
-    window.prettify_annotation = () => {
-        label_style = "style='width: calc(16% - 6px); display: inline-block; text-align:center; font-weight: bold;'";
-        input_style_base = "style='width: calc(16% - 6px); display: inline-block; padding: 5px;'";
-
-        input_style_selected = "style='width: calc(16% - 6px); display: inline-block; padding: 5px; border-width: 3px; border-color: orange; border-radius: 5px;'";
-
-        html = ""
-        for (const annotation of window.annotations) {
-            input_style = (window.hover && annotation === window.hover.annotation) ? input_style_selected : input_style_base;
-            html += `
-            <div style='margin: 0 0 20px 10px'>
-                <div>
-                    <label ${label_style}>${"id"}</label>
-                    <label ${label_style}>${"x"}</label>
-                    <label ${label_style}>${"y"}</label>
-                    <label ${label_style}>${"w"}</label>
-                    <label ${label_style}>${"h"}</label>
-                    <label ${label_style}>${"conf"}</label>
-                </div>
-                <div style='height:40px'>
-                    <input ${input_style} type='text' value='${annotation.id}'>
-                    <input ${input_style} type='text' value='${annotation.rect.startX}'>
-                    <input ${input_style} type='text' value='${annotation.rect.startY}'>
-                    <input ${input_style} type='text' value='${annotation.rect.w}'>
-                    <input ${input_style} type='text' value='${annotation.rect.h}'>
-                    <input ${input_style} type='text' value='${annotation.conf}'>
-                </div>
-            </div>`;
-        }
-        document.getElementById("annotation_display").innerHTML = html;
-    }
-    
-    window.keyup = (e) => {
-        delete keys[e.key.toLowerCase()];
-    }
-    window.keydown = (e) => {
-        console.log(e.key.toLowerCase())
-        keys[e.key.toLowerCase()] = true;
-
-        // if pressing x, delete hovered annotation
-        if (keys['x'] && window.hover) delete_annotation(window.hover.annotation);
-        if (keys['arrowright'] || keys['d']) next_frame();
-        if (keys['arrowleft'] || keys['a']) prev_frame();
-    }
-
-    window.mouse_down = (e) => {
-
-        update_mouse(e);
-
-        // If holding 'n', create new annotation
-        if (keys['n']) return create_annotation();
-
-        // else, start dragging hovered object
-        window.dragging = window.hover;
-
-        console.log(dragging)
-
-        window.draw()
-    }
-
-    window.mouse_up = () => {
-        console.log("mouseUp")
-        window.dragging = false;
-    }
-
-    window.mouse_move = (e) => {
-        ctx = window.ctx;
-        canvas = window.canvas;
-        dragging = window.dragging;
-
-        update_mouse(e);
-
-        if (!dragging) return window.draw();
-
-        annotation = dragging.annotation;
-        rect = annotation.rect
-        corner = dragging.corner;
-        
-        if (corner == window.TL) {
-            rect.w += rect.startX - mouseX;
-            rect.h += rect.startY - mouseY;
-            rect.startX = mouseX;
-            rect.startY = mouseY;
-        } else if (corner == window.TR) {
-            rect.w = mouseX - rect.startX;
-            rect.h += rect.startY - mouseY;
-            rect.startY = mouseY;
-        } else if (corner == window.BL) {
-            rect.w += rect.startX - mouseX;
-            rect.h = mouseY - rect.startY;
-            rect.startX = mouseX;
-        } else if (corner == window.BR) {
-            rect.w = mouseX - rect.startX;
-            rect.h = mouseY - rect.startY
-        }
-
-        rect.startX = Math.round(rect.startX);
-        rect.startY = Math.round(rect.startY);
-        rect.w = Math.round(rect.w);
-        rect.h = Math.round(rect.h);
-
-        // If w < 0 we have swapped sides, switch to horizontally opposite corner
-        if (rect.w < 0) {
-            rect.w = -rect.w;
-            if (corner == window.TL) corner = window.TR;
-            else if (corner == window.TR) corner = window.TL;
-            else if (corner == window.BL) corner = window.BR;
-            else if (corner == window.BR) corner = window.BL;
-        }
-        //If h < 0 we have swapped sides, switch to vertically opposite corner
-        if (rect.h < 0) {
-            rect.h = -rect.h;
-            if (corner == window.TL) corner = window.BL;
-            else if (corner == window.BL) corner = window.TL;
-            else if (corner == window.TR) corner = window.BR;
-            else if (corner == window.BR) corner = window.TR;
-        }
-        if (dragging.corner !== corner) console.log(dragging.corner + " -> " + corner);
-        dragging.corner = corner;
-
-        window.draw()
-    }
-
-    update_mouse = (e) => {
-        bodyRect = document.body.getBoundingClientRect();
-        canvasRect = e.target.getBoundingClientRect();
-        offset_x   = canvasRect.left - bodyRect.left;
-        offset_y   = canvasRect.top - bodyRect.top;
-        mouseX = e.pageX - offset_x;
-        mouseY = e.pageY - offset_y;
-
-        function sqDistance(x, y) {
-            dx = mouseX - x;
-            dy = mouseY - y;
-            return dx*dx + dy*dy;
-        }
-        
-        window.hover = false;
-        threshold = 100;
-        for (const annotation of annotations) {
-            rect = annotation.rect;
-            square_dists = [
-                sqDistance(rect.startX, rect.startY),
-                sqDistance(rect.startX + rect.w, rect.startY),
-                sqDistance(rect.startX, rect.startY + rect.h),
-                sqDistance(rect.startX + rect.w, rect.startY + rect.h),
-            ]
-
-            min_dist = Math.min(...square_dists);
-            if (min_dist > threshold) continue;
-
-            threshold = min_dist;
-            corner = square_dists.indexOf(min_dist);
-            window.hover = { corner, annotation }
-        }
-    }
-
-    create_annotation = () => {
-    
-        new_annotation = {
-            rect: {
-                startX: mouseX,
-                startY: mouseY,
-                w: 0,
-                h: 0
-            },
-            color: "rgb(255, 0, 0)",
-            id: 1,
-            conf: 1
-        };
-        annotations.push(new_annotation);
-        window.dragging = {
-            annotation: new_annotation,
-            corner: window.BL
-        }
-
-        window.draw()
-    }
-
-    delete_annotation = (annotation) => {
-        window.annotations = window.annotations.filter(function (a) {
-            return a !== annotation;
-        });
-        window.dragging = false;
-        window.hover = false;
-        window.draw();
-    }
-
-    window.init();
-
-}
-"""
-
-
-
-
-
-
-
-

annotation_editor.js → gradio_scripts/annotation_editor.js

@@ -4,8 +4,6 @@
     window.rects = [];
     window.mouseX;
     window.mouseY;
-    window.dragX;
-    window.dragY;
     window.closeEnough = 10;
     window.keys = {};
     window.hover = false;
@@ -14,7 +12,7 @@
     window.TR = 1;
     window.BL = 2;
     window.BR = 3;
-    
+
     window.frame_index = 0;
 
 
@@ -28,6 +26,10 @@
         show_frame();
     }
 
+    window.reset_annotation = () => {
+        window.frames = JSON.parse(document.getElementById("annotation_info").innerHTML);
+        show_frame();
+    }
     window.prev_frame = () => {
         window.frame_index = Math.max(window.frame_index - 1, 0);
         show_frame();
@@ -36,33 +38,35 @@
         window.frame_index = Math.min(window.frame_index + 1, window.frames.length - 1);
         show_frame();
     }
-    window.show_frame = () => {
-        const frame_info = window.frames[window.frame_index];
-        const annotations = frame_info['annotations'];
-        const frame = frame_info['frame'];
 
+    window.show_frame = () => {
+        window.frame = window.frames[window.frame_index];
 
-        window.annotations = [];
-        for (const annotation of annotations) {
-            const bbox = annotation['bbox']
-            window.annotations.push({
-                rect: {
-                    startX: bbox.left,
-                    startY: bbox.top,
-                    w: bbox.right - bbox.left,
-                    h: bbox.bottom - bbox.top
-                },
-                id: annotation.id,
-                conf: annotation.conf
-            })
-        }
+        // Load annotation from frame
+        const annotations = frame['annotations'];
+        window.annotations = annotations;
         console.log(window.annotations)
         window.dragging = false;
 
-        document.getElementById("annotation_img").src = "data:image/png;base64," + frame;
+        // Load frame image
+        const frame_img = frame['frame'];
+        document.getElementById("annotation_img").src = "data:image/png;base64," + frame_img;
+        // Draw function is called by this element using the onloaded callback
+
+        document.getElementById("annotation_frame_nbr").innerHTML = "Frame " + window.frame_index + "/" + window.frames.length;
+        
     }
 
+    // DRAW FUNCTIONS
     window.draw = () => {
+        draw_canvas();
+        draw_input_fields();
+
+        // Mark if frame is edited
+        document.getElementById("annotation_edited").style.display = (frame.edited) ? "block" : "none";
+    }
+    window.draw_canvas = () => {
+
         ctx = window.ctx;
         canvas = window.canvas;
         canvas.width = document.getElementById("annotation_img").width;
@@ -77,24 +81,24 @@
         for (const annotation of annotations) {
             //ctx.globalAlpha = annotation.conf
 
-            const rect = annotation.rect;
+            const rect = annotation.bbox;
             ctx.strokeStyle = color_from_id(annotation.id);
-            ctx.strokeRect(rect.startX, rect.startY, rect.w, rect.h);
+            ctx.strokeRect(rect.left, rect.top, rect.right - rect.left, rect.bottom - rect.top);
 
             ctx.font = "15px Arial";
             ctx.fillStyle = color_from_id(annotation.id);
             ctx.textAlign = "right";
-            ctx.fillText(annotation.id, rect.startX + rect.w, rect.startY - 3);
+            ctx.fillText(annotation.id, rect.right, rect.top - 3);
         }
 
         if (hover && !dragging) {
             annotation = hover.annotation;
-            rect = annotation.rect;
+            rect = annotation.bbox;
             handles = [
-                [rect.startX, rect.startY], 
-                [rect.startX + rect.w, rect.startY], 
-                [rect.startX, rect.startY + rect.h], 
-                [rect.startX + rect.w, rect.startY + rect.h]
+                [rect.left, rect.top], 
+                [rect.right, rect.top], 
+                [rect.left, rect.bottom], 
+                [rect.right, rect.bottom]
             ];
 
             handle = handles[hover.corner];
@@ -104,17 +108,8 @@
             ctx.rect(handle[0]-s/2, handle[1]-s/2, s, s);
             ctx.fill();
         }
-
-        prettify_annotation();
     }
-    color_from_id = (id) => {
-        //hue = Math.floor((number * 137.508 + 60) % 360)
-        power = Math.pow(2, Math.ceil(Math.log2(id))); 
-        hue = (2*id - power - 1) / power;
-        return 'hsl(' + Math.floor(hue*359) + ', 100%, 50%)'
-    }
-
-    window.prettify_annotation = () => {
+    window.draw_input_fields = () => {
         label_style = "style='width: calc(16% - 6px); display: inline-block; text-align:center; font-weight: bold;'";
         input_style_base = "style='width: calc(16% - 6px); display: inline-block; padding: 5px;'";
 
@@ -127,25 +122,34 @@
             <div style='margin: 0 0 20px 10px'>
                 <div>
                     <label ${label_style}>${"id"}</label>
-                    <label ${label_style}>${"x"}</label>
-                    <label ${label_style}>${"y"}</label>
-                    <label ${label_style}>${"w"}</label>
-                    <label ${label_style}>${"h"}</label>
+                    <label ${label_style}>${"left"}</label>
+                    <label ${label_style}>${"top"}</label>
+                    <label ${label_style}>${"right"}</label>
+                    <label ${label_style}>${"bottom"}</label>
                     <label ${label_style}>${"conf"}</label>
                 </div>
                 <div style='height:40px'>
                     <input ${input_style} type='text' value='${annotation.id}'>
-                    <input ${input_style} type='text' value='${Math.round(annotation.rect.startX)}'>
-                    <input ${input_style} type='text' value='${Math.round(annotation.rect.startY)}'>
-                    <input ${input_style} type='text' value='${Math.round(annotation.rect.w)}'>
-                    <input ${input_style} type='text' value='${Math.round(annotation.rect.h)}'>
+                    <input ${input_style} type='text' value='${Math.round(annotation.bbox.left)}'>
+                    <input ${input_style} type='text' value='${Math.round(annotation.bbox.top)}'>
+                    <input ${input_style} type='text' value='${Math.round(annotation.bbox.right)}'>
+                    <input ${input_style} type='text' value='${Math.round(annotation.bbox.bottom)}'>
                     <input ${input_style} type='text' value='${annotation.conf}'>
                 </div>
             </div>`;
         }
         document.getElementById("annotation_display").innerHTML = html;
     }
+
+    color_from_id = (id) => {
+        //hue = Math.floor((number * 137.508 + 60) % 360)
+        power = Math.pow(2, Math.ceil(Math.log2(id))); 
+        hue = (2*id - power - 1) / power;
+        return 'hsl(' + Math.floor(hue*359) + ', 100%, 50%)'
+    }
+
     
+    // KEY EVENTS
     window.keyup = (e) => {
         delete keys[e.key.toLowerCase()];
     }
@@ -159,13 +163,12 @@
         if (keys['arrowleft'] || keys['a']) prev_frame();
     }
 
+
+    // MOUSE EVENTS
     window.mouse_down = (e) => {
 
         update_mouse(e);
 
-        dragX = mouseX;
-        dragY = mouseY;
-
         // If holding 'n', create new annotation
         if (keys['n']) return create_annotation();
 
@@ -198,54 +201,43 @@
         if (!dragging) return window.draw();
 
         annotation = dragging.annotation;
-        rect = annotation.rect
+        rect = annotation.bbox;
         corner = dragging.corner;
         
-        console.log("drag", dragX, dragY)
-        console.log("d", rect.startX - mouseX, rect.startY - mouseY)
+        mouse = [mouseX, mouseY];
         if (corner == window.TL) {
-            rect.w += rect.startX - mouseX;
-            rect.h += rect.startY - mouseY;
-            rect.startX = mouseX;
-            rect.startY = mouseY;
+            [rect.left, rect.top] = mouse;
         } else if (corner == window.TR) {
-            rect.w = mouseX - rect.startX;
-            rect.h += rect.startY - mouseY;
-            rect.startY = mouseY;
+            [rect.right, rect.top] = mouse;
         } else if (corner == window.BL) {
-            rect.w += rect.startX - mouseX;
-            rect.h = mouseY - rect.startY;
-            rect.startX = mouseX;
+            [rect.left, rect.bottom] = mouse;
         } else if (corner == window.BR) {
-            rect.w = mouseX - rect.startX;
-            rect.h = mouseY - rect.startY
+            [rect.right, rect.bottom] = mouse;
         }
 
-        //rect.startX = Math.floor(rect.startX);
-        //rect.startY = Math.floor(rect.startY);
-        //rect.w = Math.floor(rect.w);
-        //rect.h = Math.floor(rect.h);
-
-        // If w < 0 we have swapped sides, switch to horizontally opposite corner
-        if (rect.w < 0) {
-            rect.w = -rect.w;
+        // If left > right we have swapped sides, switch to horizontally opposite corner
+        if (rect.left > rect.right) {
+            [rect.left, rect.right] = [rect.right, rect.left];
             if (corner == window.TL) corner = window.TR;
             else if (corner == window.TR) corner = window.TL;
             else if (corner == window.BL) corner = window.BR;
             else if (corner == window.BR) corner = window.BL;
         }
-        //If h < 0 we have swapped sides, switch to vertically opposite corner
-        if (rect.h < 0) {
-            rect.h = -rect.h;
+        //If top > bottom we have swapped sides, switch to vertically opposite corner
+        if (rect.top > rect.bottom) {
+            [rect.top, rect.bottom] = [rect.bottom, rect.top];
             if (corner == window.TL) corner = window.BL;
             else if (corner == window.BL) corner = window.TL;
             else if (corner == window.TR) corner = window.BR;
             else if (corner == window.BR) corner = window.TR;
         }
+
+        mark_frame_as_edited();
+
         if (dragging.corner !== corner) console.log(dragging.corner + " -> " + corner);
         dragging.corner = corner;
 
-        window.draw()
+        window.draw();
     }
 
     update_mouse = (e) => {
@@ -265,12 +257,12 @@
         window.hover = false;
         threshold = 100;
         for (const annotation of annotations) {
-            rect = annotation.rect;
+            rect = annotation.bbox;
             square_dists = [
-                sqDistance(rect.startX, rect.startY),
-                sqDistance(rect.startX + rect.w, rect.startY),
-                sqDistance(rect.startX, rect.startY + rect.h),
-                sqDistance(rect.startX + rect.w, rect.startY + rect.h),
+                sqDistance(rect.left, rect.top),
+                sqDistance(rect.right, rect.top),
+                sqDistance(rect.left, rect.bottom),
+                sqDistance(rect.right, rect.bottom)
             ]
 
             min_dist = Math.min(...square_dists);
@@ -282,14 +274,19 @@
         }
     }
 
+
+    // ANNOTATION UPDATES
+    mark_frame_as_edited = () => {
+        frame.edited = true;
+    }
     create_annotation = () => {
     
         new_annotation = {
-            rect: {
-                startX: mouseX,
-                startY: mouseY,
-                w: 0,
-                h: 0
+            bbox: {
+                left: mouseX,
+                top: mouseY,
+                right: mouseX,
+                bottom: mouseY
             },
             color: "rgb(255, 0, 0)",
             id: 1,
@@ -301,15 +298,19 @@
             corner: window.BL
         }
 
+        mark_frame_as_edited();
+
         window.draw()
     }
-
     delete_annotation = (annotation) => {
         window.annotations = window.annotations.filter(function (a) {
             return a !== annotation;
         });
         window.dragging = false;
         window.hover = false;
+
+        mark_frame_as_edited();
+
         window.draw();
     }
 

annotation_handler.py → gradio_scripts/annotation_handler.py

@@ -26,9 +26,6 @@ def load_frames(video, preds):
         frames = dataset.didson.load_frames(start_frame=0)
     else:
         frames = video
-
-
-    color_map = { fish['id'] : fish['color'] for fish in preds['fish'] }
     
     frame_info = []
     if len(frames):

{gradio_components → gradio_scripts}/upload_ui.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from file_reader import File
+from gradio_scripts.file_reader import File
 
 
 def Upload_Gradio(gradio_components):

inference.py

@@ -7,6 +7,7 @@ import numpy as np
 import json
 import time
 from unittest.mock import patch
+import math
 
 # assumes yolov5 on sys.path
 from lib.yolov5.models.experimental import attempt_load
@@ -24,12 +25,12 @@ WEIGHTS = 'models/v5m_896_300best.pt'
 # will need to configure these based on GPU hardware
 BATCH_SIZE = 32
 
-CONF_THRES = 0.3 # detection
-NMS_IOU  = 0.3 # NMS IOU
+CONF_THRES = 0.05 # detection
+NMS_IOU  = 0.2 # NMS IOU
+MAX_AGE = 14 # time until missing fish get's new id
+MIN_HITS = 16 # minimum number of frames with a specific fish for it to count
 MIN_LENGTH = 0.3 # minimum fish length, in meters
-MAX_AGE = 20 # time until missing fish get's new id
 IOU_THRES = 0.01 # IOU threshold for tracking
-MIN_HITS = 11 # minimum number of frames with a specific fish for it to count
 ###
 
 def norm(bbox, w, h):
@@ -60,14 +61,16 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
             inference = json_object['inference']
             width = json_object['width']
             height = json_object['height']
+            image_shapes = json_object['image_shapes']
     else:
-        inference, width, height = do_detection(dataloader, model, device, gp=gp)
+        inference, image_shapes, width, height = do_detection(dataloader, model, device, gp=gp)
 
     if save:
         json_object = {
             'inference': inference,
             'width': width,
-            'height': height
+            'height': height,
+            'image_shapes': image_shapes
         }
         json_text = json.dumps(json_object, indent=4)
         with open('static/example/inference_output.json', 'w') as f:
@@ -75,7 +78,13 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
         return
 
 
-    all_preds, real_width, real_height = do_suppression(dataloader, inference, width, height, gp=gp)
+    outputs = do_suppression(inference, gp=gp)
+
+    do_confidence_boost(inference, outputs, gp=gp)
+
+    new_outputs = do_suppression(inference, gp=gp)
+
+    all_preds, real_width, real_height = format_predictions(image_shapes, new_outputs, width, height, gp=gp)
 
     results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gp)
 
@@ -114,6 +123,7 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE, verb
     if (gp): gp(0, "Detection...")
 
     inference = []
+    image_shapes = []
     # Run detection
     with tqdm(total=len(dataloader)*batch_size, desc="Running detection", ncols=0, disable=not verbose) as pbar:
         for batch_i, (img, _, shapes) in enumerate(dataloader):
@@ -125,16 +135,24 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE, verb
             size = tuple(img.shape)
             nb, _, height, width = size  # batch size, channels, height, width
 
+
+            
             # Run model & NMS
             with torch.no_grad():
                 inf_out, _ = model(img, augment=False) 
 
+            # Save shapes for resizing to original shape
+            batch_shape = []
+            for si, pred in enumerate(inf_out):
+                batch_shape.append((img[si].shape[1:], shapes[si]))
+            image_shapes.append(batch_shape)
+            
             inference.append(inf_out)
             pbar.update(1*batch_size)
 
-    return inference, width, height
+    return inference, image_shapes, width, height
 
-def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, verbose=True):
+def do_suppression(inference, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, verbose=True):
     """
     Args:
         frames_dir: a directory containing frames to be evaluated
@@ -146,31 +164,59 @@ def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BAT
     if (gp): gp(0, "Suppression...")
     # keep predictions to feed them ordered into the Tracker
     # TODO: how to deal with large files?
-    all_preds = {}
-    with tqdm(total=len(dataloader)*batch_size, desc="Running suppression", ncols=0, disable=not verbose) as pbar:
-        for batch_i, (img, _, shapes) in enumerate(dataloader):
+    outputs = []
+    with tqdm(total=len(inference)*batch_size, desc="Running suppression", ncols=0, disable=not verbose) as pbar:
+        for batch_i, inf_out in enumerate(inference):
 
-            if gp: gp(batch_i / len(dataloader), pbar.__str__())
+            if gp: gp(batch_i / len(inference), pbar.__str__())
 
-            inf_out = inference[batch_i]
             with torch.no_grad():
                 output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres)
 
+
+            print(type(output))
+            outputs.append(output)
+
+            pbar.update(1*batch_size)
+         
+    return outputs
+
+def format_predictions(image_shapes, outputs, width, height, gp=None, batch_size=BATCH_SIZE, verbose=True):
+    """
+    Args:
+        frames_dir: a directory containing frames to be evaluated
+        image_meter_width: the width of each image, in meters (used for fish length calculation)
+        gp: a callback function which takes as input 1 parameter, (int) percent complete
+        prep_for_marking: re-index fish for manual marking output
+    """
+    
+    if (gp): gp(0, "Formatting...")
+    # keep predictions to feed them ordered into the Tracker
+    # TODO: how to deal with large files?
+    all_preds = {}
+    with tqdm(total=len(image_shapes)*batch_size, desc="Running formatting", ncols=0, disable=not verbose) as pbar:
+        for batch_i, batch in enumerate(outputs):
+
+            if gp: gp(batch_i / len(image_shapes), pbar.__str__())
+
+            batch_shapes = image_shapes[batch_i]
+
             # Format results
-            for si, pred in enumerate(output):
+            for si, pred in enumerate(batch):
+                (image_shape, original_shape) = batch_shapes[si]
                 # Clip boxes to image bounds and resize to input shape
                 clip_boxes(pred, (height, width))
                 box = pred[:, :4].clone()  # xyxy
                 confs = pred[:, 4].clone().tolist()
-                scale_boxes(img[si].shape[1:], box, shapes[si][0], shapes[si][1])  # to original shape
+                scale_boxes(image_shape, box, original_shape[0], original_shape[1])  # to original shape
                 
                 # get boxes into tracker input format - normalized xyxy with confidence score
                 # confidence score currently not used by tracker; set to 1.0
                 boxes = None
                 if box.shape[0]:
-                    real_width = shapes[si][0][1]
-                    real_height = shapes[si][0][0]
-                    do_norm = partial(norm, w=shapes[si][0][1], h=shapes[si][0][0])
+                    real_width = original_shape[0][1]
+                    real_height = original_shape[0][0]
+                    do_norm = partial(norm, w=original_shape[0][1], h=original_shape[0][0])
                     normed = list((map(do_norm, box[:, :4].tolist())))
                     boxes = np.stack([ [*bb, conf] for bb, conf in zip(normed, confs) ])
                 frame_num = (batch_i, si)
@@ -180,6 +226,67 @@ def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BAT
          
     return all_preds, real_width, real_height
 
+def do_confidence_boost(inference, safe_preds, gp=None, batch_size=BATCH_SIZE, verbose=True):
+    """
+    Args:
+        frames_dir: a directory containing frames to be evaluated
+        image_meter_width: the width of each image, in meters (used for fish length calculation)
+        gp: a callback function which takes as input 1 parameter, (int) percent complete
+        prep_for_marking: re-index fish for manual marking output
+    """
+    
+    if (gp): gp(0, "Confidence Boost...")
+    # keep predictions to feed them ordered into the Tracker
+    # TODO: how to deal with large files?
+    outputs = []
+    with tqdm(total=len(inference), desc="Running confidence boost", ncols=0, disable=not verbose) as pbar:
+        for batch_i in range(len(inference)):
+
+            if gp: gp(batch_i / len(inference), pbar.__str__())
+
+            safe = safe_preds[batch_i]
+            infer = inference[batch_i]
+
+            for i in range(len(safe)):
+                safe_frame = safe[i]
+                if len(safe_frame) == 0:
+                    continue
+
+                has_next_batch = batch_i+1 < len(inference)
+                has_prev_batch = batch_i-1 >= 0
+
+                frames = [None, None]
+                
+                next_frame = None
+                if i+1 < len(infer):
+                    next_frame = infer[i+1]
+                elif has_next_batch:
+                    next_frame = inference[batch_i + 1][0]
+
+                if next_frame != None:
+                    boost_frame(safe_frame, next_frame, 1)
+
+                prev_frame = None
+                if i-1 >= 0:
+                    prev_frame = infer[i-1]
+                elif has_prev_batch:
+                    prev_frame = inference[batch_i - 1][len(inference[batch_i - 1]) - 1]
+
+                if prev_frame != None:
+                    boost_frame(safe_frame, prev_frame, -1)
+
+            pbar.update(1*batch_size)
+          
+
+def boost_frame(safe_frame, base_frame, dt):
+    safe_boxes = safe_frame[:, :4]
+    boxes = xywh2xyxy(base_frame[:, :4])  # center_x, center_y, width, height) to (x1, y1, x2, y2)
+    ious = box_iou(boxes, safe_boxes)
+    score = torch.matmul(ious, safe_frame[:, 4])
+    # score = iou(safe_box, base_box) * confidence(safe_box)
+    base_frame[:, 4] *= 1 + (score)*math.exp(-dt*dt)
+    return base_frame
+
 def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH, verbose=True): 
 
     if (gp): gp(0, "Tracking...")

requirements.txt

@@ -1,79 +1,30 @@
-flask
-requests
-black
+# Base Requirements
+gradio
 
-matplotlib>=3.2.2
+matplotlib>=3.3
 numpy>=1.18.5
 opencv-python>=4.1.2
-Pillow
-PyYAML>=5.3.1
-scipy>=1.4.1
+Pillow>=7.1.2
 torch>=1.9.0
 torchvision>=0.8.1
 tqdm>=4.41.0
-
-tensorboard>=2.4.1
-
-seaborn>=0.11.0
 pandas
 
-thop  # FLOPs computation
+boto3 #AWS connection
 
-pycocotools
+# Requirements: lib/fish_eye ----------------------------------------------------------------
 filterpy
-celery
-redis
-boto3
-
-
-# YOLOv5 requirements
-# Usage: pip install -r requirements.txt
+scipy>=1.4.1
 
-# Base ------------------------------------------------------------------------
+# Requirements: lib/yolov5 -------------------------------------------------------------------
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch
 gitpython>=3.1.30
-matplotlib>=3.3
-numpy>=1.18.5
-opencv-python>=4.1.1
-Pillow>=7.1.2
 psutil  # system resources
-PyYAML>=5.3.1
 requests>=2.23.0
-scipy>=1.4.1
-thop>=0.1.1  # FLOPs computation
-torch>=1.7.0  # see https://pytorch.org/get-started/locally (recommended)
-torchvision>=0.8.1
-tqdm>=4.64.0
 ultralytics>=8.0.111
-# protobuf<=3.20.1  # https://github.com/ultralytics/yolov5/issues/8012
-
-# Logging ---------------------------------------------------------------------
-# tensorboard>=2.4.1
-# clearml>=1.2.0
-# comet
-
-# Plotting --------------------------------------------------------------------
-pandas>=1.1.4
 seaborn>=0.11.0
+thop>=0.1.1  # FLOPs computation
+pycocotools
+tensorboard>=2.4.1
+PyYAML>=5.3.1
 
-# Export ----------------------------------------------------------------------
-# coremltools>=6.0  # CoreML export
-# onnx>=1.10.0  # ONNX export
-# onnx-simplifier>=0.4.1  # ONNX simplifier
-# nvidia-pyindex  # TensorRT export
-# nvidia-tensorrt  # TensorRT export
-# scikit-learn<=1.1.2  # CoreML quantization
-# tensorflow>=2.4.0  # TF exports (-cpu, -aarch64, -macos)
-# tensorflowjs>=3.9.0  # TF.js export
-# openvino-dev  # OpenVINO export
-
-# Deploy ----------------------------------------------------------------------
-setuptools>=65.5.1 # Snyk vulnerability fix
-# tritonclient[all]~=2.24.0
-
-# Extras ----------------------------------------------------------------------
-# ipython  # interactive notebook
-# mss  # screenshots
-# albumentations>=1.0.3
-# pycocotools>=2.0.6  # COCO mAP

scripts/detect_frames.py

@@ -0,0 +1,138 @@
+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
+
+
+def main(args, config={}, 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
+    Args:
+        filepath (str): path to aris file
+        
+    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+    """
+    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)
+
+    dirname = args.frames
+    
+    locations = [
+        "elwha"
+    ]
+    for loc in locations:
+
+        in_loc_dir = os.path.join(dirname, loc)
+        out_dir = os.path.join(args.output, loc, "tracker", "data")
+        os.makedirs(out_dir, exist_ok=True)
+        print(in_loc_dir)
+        print(out_dir)
+
+        # run detection + tracking
+        model, device = setup_model(args.weights)
+
+        seq_list = os.listdir(in_loc_dir)
+        idx = 1
+
+        ann_list = []
+        with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
+            for seq in seq_list:
+                pbar.update(1)
+                if (seq.startswith(".")): continue
+                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)
+                frame_list = detect(in_seq_dir, out_dir, config, seq, model, device, verbose)
+                i = 0
+                for frame in frame_list:
+                    print(frame)
+                    if frame is not None:
+                        for ann in frame:
+                            print(ann)
+                            ann_list.append({
+                                'image_id': i,
+                                'category_id': 0,
+                                'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
+                                'score': ann[4]
+                            })
+                    i += 1
+    result = json.dumps(ann_list)
+    with open('../coco_eval/pred.json', 'w') as f:
+        f.write(result)
+
+                
+
+def detect(in_dir, out_dir, config, seq_name, model, device, verbose):
+    
+    #progress_log = lambda p, m: 0
+
+    # create dataloader
+    dataloader = create_dataloader_frames_only(in_dir)
+
+    try:
+        inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
+    except:
+        print("Error in " + seq_name)
+        with open(os.path.join(out_dir, "ERROR_" + seq_name + ".txt"), 'w') as f:
+            f.write("ERROR")
+        return
+
+
+    outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
+
+    frame_list = []
+    for batch_i, batch in enumerate(outputs):
+
+        batch_shapes = image_shapes[batch_i]
+
+        # Format results
+        for si, pred in enumerate(batch):
+            (image_shape, original_shape) = batch_shapes[si]
+            # Clip boxes to image bounds and resize to input shape
+            clip_boxes(pred, (height, width))
+            boxes = pred[:, :4].clone()  # xyxy
+            confs = pred[:, 4].clone().tolist()
+            scale_boxes(image_shape, boxes, original_shape[0], original_shape[1])  # to original shape
+            ann = [ [*bb, conf] for bb, conf in zip(boxes.tolist(), confs) ]
+
+            frame_list.append(ann)
+            
+    return frame_list
+
+def argument_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--frames", required=True, help="Path to frame 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")
+    return parser
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)

scripts/infer_frames.py

@@ -4,8 +4,7 @@ from datetime import datetime
 import torch
 import os
 from dataloader import create_dataloader_frames_only
-from aris import create_manual_marking, create_metadata_dictionary, prep_for_mm
-from inference import setup_model, do_suppression, do_detection, do_tracking, json_dump_round_float
+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
@@ -39,7 +38,7 @@ def main(args, config={}, verbose=True):
 
     dirname = args.frames
     
-    locations = ["kenai-val"]
+    locations = ["test"]
     for loc in locations:
 
         in_loc_dir = os.path.join(dirname, loc)
@@ -88,14 +87,21 @@ def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path, v
     dataloader = create_dataloader_frames_only(in_dir)
 
     try:
-        inference, width, height = do_detection(dataloader, model, device, verbose=verbose)
+        inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
     except:
         print("Error in " + seq_name)
         with open(os.path.join(out_dir, "ERROR_" + seq_name + ".txt"), 'w') as f:
             f.write("ERROR")
         return
 
-    all_preds, real_width, real_height = do_suppression(dataloader, inference, width, height, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
+
+    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)
 

static/example/example_result.zip

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

visualizer.py

@@ -35,9 +35,6 @@ def generate_video_batches(didson, preds, frame_rate, video_out_path, gp=None, i
     if (gp): gp(0, "Generating results video...")
     end_frame = didson.info['endframe'] or didson.info['numframes']
     out = None # need to wait til we have height and width to instantiate video file
-
-    print(preds)
-    print(preds.keys())
     
     with tqdm(total=end_frame, desc="Generating results video", ncols=0) as pbar:
         for i in range(0, end_frame, batch_size):