Spaces:

caltech-fish-counting
/

fisheye-experimental

Runtime error

App Files Files Community

oskarastrom commited on Aug 7, 2023

Commit

5ab0373

1 Parent(s): 1edac8d

Annotations

Browse files

Files changed (20) hide show

annotation_editor.js +312 -0
annotation_handler.py +69 -0
app.py +48 -2
bbox.py +309 -0
dataloader.py +56 -0
inference.py +155 -13
lib/fish_eye/sort.py +1 -1
lib/fish_eye/tracker.py +15 -5
main.py +1 -1
scripts/{infer.py → infer_aris.py} +0 -0
scripts/infer_frames.py +108 -0
state_handler.py +75 -1
static/example/2018-06-03-JD154_LeftFar_Stratum1_Set1_LO_2018-06-03_020004_3240_3840.cache +0 -0
static/example/640 2018-07-09_115439_results 2.json +0 -0
static/example/640 2018-07-09_115439_results.json +0 -0
static/example/input_conf.json +0 -0
static/example/input_conf_low.json +0 -0
static/example/input_file_vatic.txt +1540 -0
static/example/input_file_vatic.xml +0 -0
static/example/test_data/kenai-val.json +1 -0

annotation_editor.js ADDED Viewed

	@@ -0,0 +1,312 @@

+() => {
+    window.canvas = document.getElementById('canvas');
+    window.ctx = canvas.getContext('2d');
+    window.rects = [];
+    window.mouseX;
+    window.mouseY;
+    window.dragX;
+    window.dragY;
+    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='${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='${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);
+        dragX = mouseX;
+        dragY = mouseY;
+        // If holding 'n', create new annotation
+        if (keys['n']) return create_annotation();
+        // else, start dragging hovered object
+        window.dragging = false;
+        if (window.hover) {
+            window.dragging = {
+                annotation: window.hover.annotation,
+                corner: window.hover.corner
+            }
+        }
+        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;
+        console.log("drag", dragX, dragY)
+        console.log("d", rect.startX - mouseX, rect.startY - mouseY)
+        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.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 (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_handler.py ADDED Viewed

	@@ -0,0 +1,69 @@

+from dataloader import create_dataloader_aris
+from aris import BEAM_WIDTH_DIR
+import json
+import cv2
+import base64
+VIDEO_HEIGHT = 700
+def load_frames(video_path, preds_path):
+    """Load frames for annotation editing
+    """
+    dataloader, dataset = create_dataloader_aris(video_path, BEAM_WIDTH_DIR, None)
+    didson = dataset.didson
+    frames = didson.load_frames(start_frame=0)
+    frame_info, h, w = get_frame_info(frames, preds_path)
+    return frame_info
+def get_frame_info(frames, preds_path):
+    """Get visualized video frames ready for output, given raw ARIS/DIDSON frames.
+    Warning: all frames in frames will be stored in memory - careful of OOM errors. Consider processing large files
+    in batches, such as in generate_video_batches()
+    Returns:
+        list(np.ndarray), height (int), width (int)
+    """
+    preds = json.load(open(preds_path, 'r'))
+    color_map = { fish['id'] : fish['color'] for fish in preds['fish'] }
+    frame_info = []
+    if len(frames):
+        # assumes all frames the same size
+        h, w = frames[0].shape
+        # enforce a standard size so that text/box thickness is consistent
+        scale_factor = VIDEO_HEIGHT / h
+        h = VIDEO_HEIGHT
+        w = int(scale_factor*w)
+        num_frames = min(len(frames), len(preds['frames']))
+        for i, frame_raw in enumerate(frames[:num_frames]):
+            image = cv2.resize(cv2.cvtColor(frame_raw, cv2.COLOR_GRAY2BGR), (w,h))
+            retval, buffer = cv2.imencode('.jpg', image)
+            jpg_as_text = base64.b64encode(buffer).decode("utf-8")
+            frame = {
+                'annotations': [],
+                'frame': jpg_as_text
+            }
+            for fish in preds['frames'][i]['fish']:
+                xmin, ymin, xmax, ymax = fish['bbox']
+                hexx = color_map[fish['fish_id']].lstrip('#')
+                frame['annotations'].append({
+                    'bbox': {
+                        'left': int(round(xmin * w)),
+                        'right': int(round(xmax * w)),
+                        'top': int(round(ymin * h)),
+                        'bottom': int(round(ymax * h)),
+                    },
+                    'id': str(fish['fish_id']),
+                    'conf': fish['conf']
+                })
+            frame_info.append(frame)
+    return frame_info, h, w

app.py CHANGED Viewed

@@ -7,7 +7,12 @@ import numpy as np
 from aws_handler import upload_file
 from aris import create_metadata_table
 import cv2
 table_headers = ["TOTAL" , "FRAME_NUM", "DIR", "R", "THETA", "L", "TIME", "DATE", "SPECIES"]
 info_headers = [
     "TOTAL_TIME", "DATE", "START", "END",
@@ -26,6 +31,9 @@ css = """
 #marking_json label {
     height: calc(100% - 30px) !important;
 }
 """
 js_update_tabs = """
     async () => {
@@ -50,6 +58,7 @@ result = {}
 # Start function, called on file upload
 def on_input(file_list):
@@ -78,6 +87,7 @@ def handle_next(_, progress=gr.Progress()):
         set_progress = lambda pct, msg : progress(pct, desc="File " + str(state['index']+1) + "/" + str(state['total']) + ": " + msg)
     set_progress(0, "Starting...")
     file_info = state['files'][state['index']]
     file_name = file_info[0].split("/")[-1]
     bytes = file_info[1]
@@ -85,14 +95,20 @@ def handle_next(_, progress=gr.Progress()):
     print(dir_name)
     print(file_path)
     if not valid:
         return {
             result_handler: gr.update(),
             inference_handler: gr.update()
         }
     upload_file(file_path, "fishcounting", "webapp_uploads/" + file_name)
     metadata, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(file_path, gradio_progress=set_progress)
     result["path_video"].append(video_filepath)
     result["path_zip"].append(zip_filepath)
     result["path_json"].append(json_filepath)
@@ -101,8 +117,11 @@ def handle_next(_, progress=gr.Progress()):
     result["fish_table"].append(fish_table)
     result["fish_info"].append(fish_info)
     state['index'] += 1
     return {
         result_handler: gr.update(value = str(state["index"])),
         inference_handler: gr.update()
@@ -115,16 +134,19 @@ def show_example_data():
     return gr.update(value=str(state["index"]))
 def show_data():
     i = state["index"] - 1
-    # Only show result for up to max_tabs files
     if i >= max_tabs:
         return {
             zip_out: gr.update(value=result["path_zip"])
         }
     not_done = state['index'] < state['total']
     return {
         zip_out: gr.update(value=result["path_zip"]),
         tabs[i]['tab']: gr.update(),
@@ -137,7 +159,22 @@ def show_data():
-max_tabs = 10
 demo = gr.Blocks()
 with demo:
     with gr.Blocks(css=css) as inner_body:
@@ -150,6 +187,15 @@ with demo:
         <style id="tab_style"></style>
         """
         )
         #Input field for aris submission
         input = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")

 from 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
+import json
+max_tabs = 10
 table_headers = ["TOTAL" , "FRAME_NUM", "DIR", "R", "THETA", "L", "TIME", "DATE", "SPECIES"]
 info_headers = [
     "TOTAL_TIME", "DATE", "START", "END",
 #marking_json label {
     height: calc(100% - 30px) !important;
 }
+#canvas {
+    align-self: center;
+}
 """
 js_update_tabs = """
     async () => {
 # Start function, called on file upload
 def on_input(file_list):
         set_progress = lambda pct, msg : progress(pct, desc="File " + str(state['index']+1) + "/" + str(state['total']) + ": " + msg)
     set_progress(0, "Starting...")
+    # Save file and create a new directory for result
     file_info = state['files'][state['index']]
     file_name = file_info[0].split("/")[-1]
     bytes = file_info[1]
     print(dir_name)
     print(file_path)
+    # Check that the file was valid
     if not valid:
         return {
             result_handler: gr.update(),
             inference_handler: gr.update()
         }
+    # Send uploaded file to AWS
     upload_file(file_path, "fishcounting", "webapp_uploads/" + file_name)
+    # Do inference
     metadata, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(file_path, gradio_progress=set_progress)
+    # Store result for that file
     result["path_video"].append(video_filepath)
     result["path_zip"].append(zip_filepath)
     result["path_json"].append(json_filepath)
     result["fish_table"].append(fish_table)
     result["fish_info"].append(fish_info)
+    # Increase file index
     state['index'] += 1
+    # Send of update to result_handler to show new result
+    # Leave inference_handler update blank to avoid starting next inference until result is updated
     return {
         result_handler: gr.update(value = str(state["index"])),
         inference_handler: gr.update()
     return gr.update(value=str(state["index"]))
 def show_data():
+    # Get last index
     i = state["index"] - 1
+    # If index is larger than max_tabs, only add file to zip list
     if i >= max_tabs:
         return {
             zip_out: gr.update(value=result["path_zip"])
         }
+    # Check if inference is done
     not_done = state['index'] < state['total']
+    # Send update to UI, and to inference_handler to start next file inference
     return {
         zip_out: gr.update(value=result["path_zip"]),
         tabs[i]['tab']: gr.update(),
+def init_annotation():
+    frame_info = load_frames("static/example/input_file.aris", "static/example/640 2018-07-09_115439_results.json")
+    html = "<div style='display:flex'>"
+    html += "<canvas id='canvas' style='width:50%' onmousedown='mouse_down(event)' onmousemove='mouse_move(event)' onmouseup='mouse_up()' onmouseleave='mouse_up()'></canvas>"
+    html += "<div id='annotation_display' style='width:50%'></div>"
+    html += "</div>"
+    html += "<p id='annotation_info' style='display:none'>" + json.dumps(frame_info) + "</p>"
+    html += "<img id='annotation_img' onload='draw()' style='display:none'></img>"
+    return html
 demo = gr.Blocks()
 with demo:
     with gr.Blocks(css=css) as inner_body:
         <style id="tab_style"></style>
         """
         )
+        annotation_editor = gr.HTML("""""")
+        with open('annotation_editor.js', 'r') as f:
+            js = f.read()
+            print(js)
+            annotation_editor.change(lambda x: gr.update(), None, annotation_editor, _js=js)
+        gr.Button("Edit Annotation").click(init_annotation, None, annotation_editor)
         #Input field for aris submission
         input = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")

bbox.py ADDED Viewed

	@@ -0,0 +1,309 @@

+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();
+}
+"""

dataloader.py CHANGED Viewed

@@ -9,6 +9,7 @@ import struct
 from contextlib import contextmanager
 import torch
 from torch.utils.data import Dataset
 # assumes yolov5 on sys.path
 from lib.yolov5.utils.general import xyxy2xywh
@@ -69,11 +70,66 @@ def create_dataloader_frames(frames_path, batch_size=32, model_stride_max=32,
                                        hyp=None, cache=None, rect=True, rank=rank,
                                        workers=workers, pad=pad)[0]
 # # # # # #
 # End factory(ish) methods
 # # # # # #
 class ARISBatchedDataset(Dataset):
     def __init__(self, aris_filepath, beam_width_dir, annotations_file, batch_size, num_frames_bg_subtract=1000, disable_output=False,
                     cache_bg_frames=False):

 from contextlib import contextmanager
 import torch
 from torch.utils.data import Dataset
+import torchvision.transforms as T
 # assumes yolov5 on sys.path
 from lib.yolov5.utils.general import xyxy2xywh
                                        hyp=None, cache=None, rect=True, rank=rank,
                                        workers=workers, pad=pad)[0]
+def create_dataloader_frames_only(frames_path, batch_size=32, img_size=896, workers=0):
+    """
+    Create a DataLoader for a directory of frames without labels.
+    Args:
+        model_stride_max: use model.stride.max()
+    """
+    return YOLOFrameDataset(frames_path, img_size=img_size, batch_size=batch_size)
 # # # # # #
 # End factory(ish) methods
 # # # # # #
+import os
+import pandas as pd
+from torchvision.io import read_image
+import re
+class YOLOFrameDataset(Dataset):
+    def __init__(self, img_dir, img_size=896, batch_size=32):
+        self.img_dir = img_dir
+        self.img_size = img_size
+        self.files = os.listdir(img_dir)
+        self.files = list(filter(lambda f: f[-4:] == ".jpg", self.files))
+        self.files.sort(key=lambda f: int(re.sub('\D', '', f)))
+        print(self.files)
+        n = len(self.files)
+        self.batches = []
+        for i in range(0,n,batch_size):
+            self.batches.append((i, min(n, i+batch_size)))
+        print(self.batches)
+    def __len__(self):
+        return len(self.batches)
+    def __iter__(self):
+        for batch_idx in self.batches:
+            batch = []
+            shapes = []
+            for i in range(batch_idx[0], batch_idx[1]):
+                img_name = self.files[i]
+                img_path = os.path.join(self.img_dir, img_name)
+                img = read_image(img_path)
+                shapes.append([[img.shape[1], img.shape[2]], None])
+                ratio = self.img_size / max(img.shape[1], img.shape[2])
+                transform = T.Resize([int(ratio*img.shape[1]), int(ratio*img.shape[2])])
+                img = transform(img)
+                batch.append(img)
+            image = torch.stack(batch)
+            yield (image, None, shapes)
 class ARISBatchedDataset(Dataset):
     def __init__(self, aris_filepath, beam_width_dir, annotations_file, batch_size, num_frames_bg_subtract=1000, disable_output=False,
                     cache_bg_frames=False):

inference.py CHANGED Viewed

@@ -5,23 +5,27 @@ from tqdm import tqdm
 from functools import partial
 import numpy as np
 import json
 from unittest.mock import patch
 # 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 non_max_suppression
-from lib.yolov5.utils.general import clip_boxes, scale_boxes
 from lib.fish_eye.tracker import Tracker
 ### Configuration options
 WEIGHTS = 'models/v5m_896_300best.pt'
 # will need to configure these based on GPU hardware
 BATCH_SIZE = 32
-conf_thres = 0.3 # detection
-iou_thres  = 0.3 # NMS IOU
 min_length = 0.3 # minimum fish length, in meters
 ###
@@ -44,7 +48,31 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
     model, device = setup_model(weights)
-    all_preds = do_detection(dataloader, model, device, gp=gp)
     results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gp)
@@ -63,7 +91,7 @@ def setup_model(weights_fp=WEIGHTS, imgsz=896, batch_size=32):
     half = device.type != 'cpu'  # half precision only supported on CUDA
     if half:
         model.half()
-    model.eval();
     # Create dataloader for batched inference
     img = torch.zeros((1, 3, imgsz, imgsz), device=device)
@@ -81,22 +109,49 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE):
     """
     if (gp): gp(0, "Detection...")
-    # keep predictions to feed them ordered into the Tracker
-    # TODO: how to deal with large files?
-    all_preds = {}
     # Run detection
     with tqdm(total=len(dataloader)*batch_size, desc="Running detection", ncols=0) as pbar:
         for batch_i, (img, _, shapes) in enumerate(dataloader):
             if gp: gp(batch_i / len(dataloader), pbar.__str__())
             img = img.to(device, non_blocking=True)
             img = img.half() if device.type != 'cpu' else img.float()  # uint8 to fp16/32
             img /= 255.0  # 0 - 255 to 0.0 - 1.0
-            nb, _, height, width = img.shape  # batch size, channels, height, width
             # Run model & NMS
             with torch.no_grad():
                 inf_out, _ = model(img, augment=False)
                 output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres)
             # Format results
@@ -115,6 +170,7 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE):
                     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)
                 all_preds[frame_num] = boxes
             pbar.update(1*batch_size)
@@ -132,7 +188,7 @@ def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None):
         'image_meter_width': image_meter_width,
         'image_meter_height': image_meter_height
     }
-    tracker = Tracker(clip_info, args={ 'max_age': 9, 'min_hits': 0, 'iou_threshold': 0.01}, min_hits=11)
     # Run tracking
     with tqdm(total=len(all_preds), desc="Running tracking", ncols=0) as pbar:
@@ -144,6 +200,7 @@ def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None):
             else:
                 tracker.update()
             pbar.update(1)
     json_data = tracker.finalize(min_length=min_length)
     return json_data
@@ -164,4 +221,89 @@ def json_dump_round_float(some_object, out_path, num_digits=4):
         return of(*args, **kwargs)
     with patch('json.encoder._make_iterencode', wraps=inner):
-        return json.dump(some_object, open(out_path, 'w'), indent=2)

 from functools import partial
 import numpy as np
 import json
+import time
 from unittest.mock import patch
 # 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 lib.fish_eye.tracker import Tracker
 ### Configuration options
 WEIGHTS = 'models/v5m_896_300best.pt'
 # will need to configure these based on GPU hardware
 BATCH_SIZE = 32
+conf_thres = 0.001 # detection
+iou_thres  = 0.1 # NMS IOU
 min_length = 0.3 # minimum fish length, in meters
 ###
     model, device = setup_model(weights)
+    load = False
+    save = False
+    if load:
+        with open('static/example/inference_output.json', 'r') as f:
+            json_object = json.load(f)
+            inference = json_object['inference']
+            width = json_object['width']
+            height = json_object['height']
+    else:
+        inference, width, height = do_detection(dataloader, model, device, gp=gp)
+    if save:
+        json_object = {
+            'inference': inference,
+            'width': width,
+            'height': height
+        }
+        json_text = json.dumps(json_object, indent=4)
+        with open('static/example/inference_output.json', 'w') as f:
+            f.write(json_text)
+        return
+    all_preds = do_suppression(dataloader, inference, width, height, gp=gp)
     results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gp)
     half = device.type != 'cpu'  # half precision only supported on CUDA
     if half:
         model.half()
+    model.eval()
     # Create dataloader for batched inference
     img = torch.zeros((1, 3, imgsz, imgsz), device=device)
     """
     if (gp): gp(0, "Detection...")
+    inference = []
     # Run detection
     with tqdm(total=len(dataloader)*batch_size, desc="Running detection", ncols=0) as pbar:
         for batch_i, (img, _, shapes) in enumerate(dataloader):
             if gp: gp(batch_i / len(dataloader), pbar.__str__())
             img = img.to(device, non_blocking=True)
             img = img.half() if device.type != 'cpu' else img.float()  # uint8 to fp16/32
             img /= 255.0  # 0 - 255 to 0.0 - 1.0
+            size = tuple(img.shape)
+            nb, _, height, width = size  # batch size, channels, height, width
+            print(nb, _, height, width)
             # Run model & NMS
             with torch.no_grad():
                 inf_out, _ = model(img, augment=False)
+            inference.append(inf_out)
+            pbar.update(1*batch_size)
+    return inference, width, height
+def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BATCH_SIZE):
+    """
+    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, "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) as pbar:
+        for batch_i, (img, _, shapes) in enumerate(dataloader):
+            if gp: gp(batch_i / len(dataloader), 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)
             # Format results
                     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)
+                print(frame_num)
                 all_preds[frame_num] = boxes
             pbar.update(1*batch_size)
         'image_meter_width': image_meter_width,
         'image_meter_height': image_meter_height
     }
+    tracker = Tracker(clip_info, args={ 'max_age': 20, 'min_hits': 0, 'iou_threshold': 0.01}, min_hits=11)
     # Run tracking
     with tqdm(total=len(all_preds), desc="Running tracking", ncols=0) as pbar:
             else:
                 tracker.update()
             pbar.update(1)
     json_data = tracker.finalize(min_length=min_length)
     return json_data
         return of(*args, **kwargs)
     with patch('json.encoder._make_iterencode', wraps=inner):
+        return json.dump(some_object, open(out_path, 'w'), indent=2)
+def non_max_suppression(
+        prediction,
+        conf_thres=0.25,
+        iou_thres=0.45,
+        max_det=300,
+):
+    """Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
+    Returns:
+         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
+    """
+    # Checks
+    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+    if isinstance(prediction, (list, tuple)):  # YOLOv5 model in validation model, output = (inference_out, loss_out)
+        prediction = prediction[0]  # select only inference output
+    device = prediction.device
+    mps = 'mps' in device.type  # Apple MPS
+    if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
+        prediction = prediction.cpu()
+    bs = prediction.shape[0]  # batch size
+    xc = prediction[..., 4] > conf_thres  # candidates
+    # Settings
+    # min_wh = 2  # (pixels) minimum box width and height
+    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
+    redundant = True  # require redundant detections
+    merge = False  # use merge-NMS
+    output = [torch.zeros((0, 6), device=prediction.device)] * bs
+    for xi, x in enumerate(prediction):  # image index, image inference
+        # Keep boxes that pass confidence threshold
+        x = x[xc[xi]]  # confidence
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+        # Compute conf
+        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+        # Box/Mask
+        box = xywh2xyxy(x[:, :4])  # center_x, center_y, width, height) to (x1, y1, x2, y2)
+        mask = x[:, 6:]  # zero columns if no masks
+        # Detections matrix nx6 (xyxy, conf, cls)
+        conf, j = x[:, 5:6].max(1, keepdim=True)
+        x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
+        # Check shape
+        n = x.shape[0]  # number of boxes
+        if not n:  # no boxes
+            continue
+        x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence and remove excess boxes
+        # Batched NMS
+        boxes  = x[:, :4]  # boxes (offset by class), scores
+        scores = x[:, 4]
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
+        i = i[:max_det]  # limit detections
+        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
+            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
+            weights = iou * scores[None]  # box weights
+            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+            if redundant:
+                i = i[iou.sum(1) > 1]  # require redundancy
+        output[xi] = x[i]
+        if mps:
+            output[xi] = output[xi].to(device)
+        logging = False
+    return output

lib/fish_eye/sort.py CHANGED Viewed

@@ -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)
         # update matched trackers with assigned detections
         for m in matched:

         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)
         # update matched trackers with assigned detections
         for m in matched:

lib/fish_eye/tracker.py CHANGED Viewed

@@ -20,12 +20,20 @@ class Tracker:
     def update(self, dets=np.empty((0, 5))):
         new_frame_entries = []
         for track in self.algorithm.update(dets):
             self.fish_ids[int(track[4])] += 1
             new_frame_entries.append({
-            'fish_id': int(track[4]),
-            'bbox': list(track[:4]),
-            'visible': 1,
-            'human_labeled': 0
             })
         new_frame_entries = sorted(new_frame_entries, key=lambda k: k['fish_id'])
@@ -38,7 +46,7 @@ class Tracker:
     def finalize(self, output_path=None, min_length=-1.0): # vert_margin=0.0
         json_data = deepcopy(self.json_data)
         # map (valid) fish IDs to 0, 1, 2, ...
         fish_id_map = {}
         for fish_id, count in self.fish_ids.items():
@@ -95,6 +103,8 @@ class Tracker:
             json_data['fish'].append(fish_entry)
         # filter 'fish' field by fish length
         json_data = Fish_Length.add_lengths(json_data)
         invalid_ids = []

     def update(self, dets=np.empty((0, 5))):
         new_frame_entries = []
         for track in self.algorithm.update(dets):
+            conf = 0
+            min_score = 1000000
+            for det in dets:
+                score = sum(abs(det[0:4] - track[0:4]))
+                if (score < min_score):
+                    min_score = score
+                    conf = det[4]
             self.fish_ids[int(track[4])] += 1
             new_frame_entries.append({
+                'fish_id': int(track[4]),
+                'bbox': list(track[:4]),
+                'visible': 1,
+                'human_labeled': 0,
+                'conf': conf
             })
         new_frame_entries = sorted(new_frame_entries, key=lambda k: k['fish_id'])
     def finalize(self, output_path=None, min_length=-1.0): # vert_margin=0.0
         json_data = deepcopy(self.json_data)
         # map (valid) fish IDs to 0, 1, 2, ...
         fish_id_map = {}
         for fish_id, count in self.fish_ids.items():
             json_data['fish'].append(fish_entry)
         # filter 'fish' field by fish length
         json_data = Fish_Length.add_lengths(json_data)
         invalid_ids = []

main.py CHANGED Viewed

@@ -48,7 +48,7 @@ def predict_task(filepath, weights=WEIGHTS, gradio_progress=None):
     frame_rate = dataset.didson.info['framerate']
     # run detection + tracking
-    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, weights=WEIGHTS)
     # re-index results if desired - this should be done before writing the file
     results = prep_for_mm(results)

     frame_rate = dataset.didson.info['framerate']
     # run detection + tracking
+    results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, weights=weights)
     # re-index results if desired - this should be done before writing the file
     results = prep_for_mm(results)

scripts/{infer.py → infer_aris.py} RENAMED Viewed

File without changes

scripts/infer_frames.py ADDED Viewed

	@@ -0,0 +1,108 @@

+import project_path
+import argparse
+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 visualizer import generate_video_batches
+import json
+from tracker import Tracker
+def main(args):
+    """
+    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())
+    gradio_progress = lambda p, m: print(m, p)
+    dirname = args.frames
+    print(dirname)
+    filename = os.path.basename(dirname.split("/")[-1])
+    print(filename)
+    results_filepath = os.path.join(dirname, f"{filename}_results.json")
+    os.makedirs(dirname, exist_ok=True)
+    image_meter_width = -1
+    image_meter_height = -1
+    with open(args.metadata, 'r') as f:
+        json_object = json.loads(f.read())
+        print(len(json_object))
+        for seq in json_object:
+            if seq['clip_name'] == filename:
+                print(seq)
+                image_meter_width = seq['x_meter_stop'] - seq['x_meter_start']
+                image_meter_height = seq['y_meter_stop'] - seq['y_meter_start']
+    if (image_meter_height == -1):
+        print("No metadata found for file " + filename)
+        return
+    # create dataloader
+    print("Initializing Dataloader...")
+    dataloader = create_dataloader_frames_only(dirname)
+    # extract aris/didson info. didson does not yet have pixel-meter info
+    # TODO: load from metadata
+    # run detection + tracking
+    model, device = setup_model(args.weights)
+    inference, width, height = do_detection(dataloader, model, device, gp=gradio_progress)
+    print(len(inference[0]), len(inference[1]))
+    all_preds = do_suppression(dataloader, inference, width, height, gp=gradio_progress)
+    print(len(all_preds.keys()))
+    results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gradio_progress)
+    print(results)
+    print(str(width), str(height))
+    mot_rows = []
+    for frame in results['frames']:
+        for fish in frame['fish']:
+            bbox = fish['bbox']
+            row = []
+            row.append(str(frame['frame_num']))
+            row.append(str(fish['fish_id'] + 1))
+            row.append(str(int(bbox[0]*width)))
+            row.append(str(int(bbox[1]*height)))
+            row.append(str(int(bbox[2]*width)))
+            row.append(str(int(bbox[3]*height)))
+            row.append("-1")
+            row.append("-1")
+            row.append("-1")
+            row.append("-1")
+            mot_rows.append(",".join(row))
+    mot_text = "\n".join(mot_rows)
+    print(mot_text)
+    with open(os.path.join(dirname, "gt.txt"), 'w') as f:
+        f.write(mot_text)
+    return
+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 metadata file. 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)

state_handler.py CHANGED Viewed

@@ -1,4 +1,5 @@
 from aris import create_metadata_table
 example_metadata = {
     "FILE_NAME":  "static/example_metadata/fisheye",
@@ -372,4 +373,77 @@ def reset_state(result, state):
     # Reset State
     state['files'] = []
     state['index'] = 0
-    state['total'] = 0

 from aris import create_metadata_table
+import json
 example_metadata = {
     "FILE_NAME":  "static/example_metadata/fisheye",
     # 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")

static/example/2018-06-03-JD154_LeftFar_Stratum1_Set1_LO_2018-06-03_020004_3240_3840.cache ADDED Viewed

Binary file (11.1 kB). View file

static/example/640 2018-07-09_115439_results 2.json ADDED Viewed