Batch Annotation Loading

app.py

@@ -11,14 +11,16 @@ from zipfile import ZipFile
 import os
 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
-
+from dataloader import create_dataloader_aris
+from aris import BEAM_WIDTH_DIR
 
 #Initialize State & Result
 state = {
     'files': [],
     'index': 1,
     'total': 1,
-    'annotation_index': -1
+    'annotation_index': -1,
+    'frame_index': 0
 }
 result = {}
 
@@ -180,40 +182,71 @@ def cancel_inference():
 def prepare_annotation(index):
 
     state['annotation_index'] = index
+    state['frame_index'] = 0
 
+    if result["aris_input"][index]:
+        return {
+            annotation_progress: gr.update(value="<p id='annotation_info' style='display:none'>[]</p><!--" + str(np.random.rand()) + "-->", visible=True),
+            master_tabs: gr.update(selected=2)
+        }
     return {
-        annotation_progress: gr.update(value="<p align='center' style='font-size: large;font-style: italic;'>Loading annotation...</p><!--" + str(np.random.rand()) + "-->", visible=True),
-        master_tabs: gr.update(selected=2)
+        annotation_progress: gr.update(),
+        master_tabs: gr.update()
     }
 
-# Load frames and annotation information and show
-def open_annotation(_, progress=gr.Progress()):
+annotation_info = None
+annotation_dataset = None
+# annotation_progress.change
+def load_annotation(_, progress=gr.Progress()):
+    global annotation_info, annotation_dataset
+
+    # Get result index
     result_index = state['annotation_index']
 
     set_progress = lambda pct, msg: progress(pct, desc=msg)
 
+    if state['frame_index'] == 0:
+        if set_progress: set_progress(0, "Loading Frames")
+        dataloader, annotation_dataset = create_dataloader_aris(result["aris_input"][result_index], BEAM_WIDTH_DIR, None)
+
+    # Check that frames remain to be loaded
+    if state['frame_index'] < len(result['json_result'][result_index]['frames']):
+
+        # load frames and annotation
+        annotation_info, state['frame_index'] = init_frames(annotation_dataset, result['json_result'][result_index], state['frame_index'], gp=set_progress)
+    
+        # save as html element
+        annotation_storage = "<p id='annotation_info' style='display:none'>" + json.dumps(annotation_info) + "</p>"
+
+        return {
+            annotation_editor: gr.update(), 
+            annotation_progress: gr.update(value=annotation_storage)
+        }
+
+    # If complete, start annotation editor    
+
     annotation_html = ""
-    if result["aris_input"][result_index]:
-        annotations = init_frames(result["aris_input"][result_index], result['json_result'][result_index], gp=set_progress)
-
-        # 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(annotations) + "</p>"
-        annotation_html += "<img id='annotation_img' onload='draw()' style='display:none'></img>"
-        annotation_html += "<!--" + str(np.random.rand()) + "-->"
 
-    return gr.update(value=annotation_html, visible=True), gr.update(visible=False), gr.update(visible=True)
+    # Header
+    annotation_html += "<div id='annotation_header'>"
+    annotation_html += "     <h1 id='annotation_frame_nbr'>Frame 0/100</h1>"
+    annotation_html += "     <p id='annotation_edited'>(edited)</p>"
+    annotation_html += "</div>"
+
+    # Annotation Body
+    annotation_html += "<div style='display:flex'>"
+    annotation_html += "     <canvas id='canvas' style='width:50%' onmousedown='mouse_down(event)' onmousemove='mouse_move(event)' onmouseup='mouse_up()' onmouseleave='mouse_up()'></canvas>"
+    annotation_html += "     <div id='annotation_display' style='width:50%'></div>"
+    annotation_html += "</div>"
+    
+    # Dummy objects
+    annotation_html += "<img id='annotation_img' onload='draw()' style='display:none'></img>"
+    annotation_html += "<!--" + str(np.random.rand()) + "-->"
+
+    return {
+        annotation_editor: gr.update(value=annotation_html, visible=True),
+        annotation_progress: gr.update(visible=False)
+    }
 
 
 components = {}
@@ -290,8 +323,23 @@ with demo:
                 # Add annotation editor component
                 annotation_editor = gr.HTML("", visible=False)
 
+                annotation_storage = gr.HTML("", visible=False)
+
                 # Event listener for opening annotation 
-                annotation_progress.change(open_annotation, annotation_progress, [annotation_editor, annotation_progress])
+                annotation_progress.change(load_annotation, annotation_progress, [annotation_editor, annotation_progress], _js="""
+                    () => {
+                        info_string = document.getElementById("annotation_info").innerHTML;
+                        info = JSON.parse(info_string);
+                        console.log(info)
+                        if (info.length == 0) {
+                            window.annotation_info = [];
+                            return false;
+                        } 
+                        window.annotation_info = window.annotation_info.concat(info)
+                        console.log(window.annotation_info)
+                        return true;
+                    }
+                """)
                 
                 # Event listener for running javascript defined in 'annotation_editor.js'
                 with open('gradio_scripts/annotation_editor.js', 'r') as f:

gradio_scripts/annotation_editor.js

@@ -17,7 +17,8 @@
 
 
     window.init = () => {
-        window.frames = JSON.parse(document.getElementById("annotation_info").innerHTML);
+        window.frames = window.annotation_info;
+        console.log(window.frames);
         window.frame_index = 0;
 
         document.removeEventListener('keydown', keydown);
@@ -51,8 +52,8 @@
         window.dragging = false;
 
         // Load frame image
-        const frame_img = frame['frame'];
-        document.getElementById("annotation_img").src = "annotation_frame_dir/" + frame_index + ".jpg";
+        const frame_img = frame['base64'];
+        document.getElementById("annotation_img").src = "data:image/jpeg;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;

gradio_scripts/annotation_handler.py

@@ -1,5 +1,3 @@
-from dataloader import create_dataloader_aris
-from aris import BEAM_WIDTH_DIR
 import json
 import cv2
 import base64
@@ -7,7 +5,7 @@ import base64
 VIDEO_HEIGHT = 700
 
 
-def init_frames(video, preds, gp=None):
+def init_frames(dataset, preds, index, gp=None):
     """Load frames for annotation editing
 
     
@@ -21,10 +19,6 @@ def init_frames(video, preds, gp=None):
                 )
             })
     """
-
-    if gp: gp(0, "Loading Frames")
-
-    dataloader, dataset = create_dataloader_aris(video, BEAM_WIDTH_DIR, None)
     images = dataset.didson.load_frames(start_frame=0, end_frame=1)
 
     # assumes all frames the same size
@@ -40,17 +34,21 @@ def init_frames(video, preds, gp=None):
     if gp: gp(0, "Extracting Frames")
     if len(preds['frames']):
         
-        for i, frame_info in enumerate(preds['frames']):
-            if gp: gp(i/len(preds['frames']), "Extracting Frames")
+        end_index = min(index+1000, len(preds['frames']))
+        for i, frame_info in enumerate(preds['frames'][index:end_index]):
+            if gp: gp((index + i)/len(preds['frames']), "Extracting Frames")
 
             # Extract frames
             img_raw = dataset.didson.load_frames(start_frame=i, end_frame=i+1)[0]
             image = cv2.resize(cv2.cvtColor(img_raw, cv2.COLOR_GRAY2BGR), (w, h))
-            cv2.imwrite("annotation_frame_dir/" + str(i) + ".jpg", image)
+            #cv2.imwrite("annotation_frame_dir/" + str(i) + ".jpg", image)
+            retval, buffer = cv2.imencode('.jpg', image)
+            b64 = base64.b64encode(buffer).decode("utf-8")
 
             # Extract annotations
             frame = {
                 'annotations': [],
+                'base64': b64
             }
             for fish in frame_info['fish']:
                 xmin, ymin, xmax, ymax = fish['bbox']
@@ -66,5 +64,5 @@ def init_frames(video, preds, gp=None):
                 })
             annotations.append(frame)
     
-    return annotations
+    return annotations, end_index