ByteTrack for UI

app.py

@@ -14,6 +14,8 @@ from gradio_scripts.result_ui import Result_Gradio, update_result, table_headers
 from dataloader import create_dataloader_aris
 from aris import BEAM_WIDTH_DIR
 
+WEBAPP_VERSION = "1.0"
+
 #Initialize State & Result
 state = {
     'files': [],
@@ -27,24 +29,29 @@ result = {}
 
 
 # Called when an Aris file is uploaded for inference
-def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_age, use_associative, boost_power, boost_decay, min_length, min_travel):
+def on_aris_input(file_list, model_id, conf_thresh, iou_thresh, min_hits, max_age, associative_tracker, boost_power, boost_decay, byte_low_conf, byte_high_conf, min_length, min_travel):
 
     # Reset Result
     reset_state(result, state)
     state['files'] = file_list
     state['total'] = len(file_list)
+    state['version'] = WEBAPP_VERSION
     state['hyperparams'] = {
         'model': models[model_id] if model_id in models else models['master'],
         'conf_thresh': conf_thresh,
         'iou_thresh': iou_thresh,
         'min_hits': min_hits,
         'max_age': max_age,
-        'use_associative_tracking': use_associative,
-        'boost_power': boost_power,
-        'boost_decay': boost_decay,
         'min_length': min_length,
-        'min_travel': min_travel
+        'min_travel': min_travel,
+        'associative_tracker': associative_tracker,
     }
+    if (associative_tracker == "Confidence Boost"):
+        state['hyperparams']['boost_power'] = boost_power
+        state['hyperparams']['boost_decay'] = boost_decay
+    elif (associative_tracker == "ByteTrack"):
+        state['hyperparams']['byte_low_conf'] = byte_low_conf
+        state['hyperparams']['byte_high_conf'] = byte_high_conf
 
     print(" ")
     print("Running with:")
@@ -69,6 +76,7 @@ def on_result_upload(zip_list, aris_list):
     
 
     reset_state(result, state)
+    state['version'] = WEBAPP_VERSION
 
     component_updates = { 
         master_tabs: gr.update(selected=1),

aris.py

@@ -441,6 +441,15 @@ def create_metadata_table(result, table_headers, info_headers):
     else:
         metadata = { 'FISH': [] }
     
+    # Calculate detection dropout
+    for fish in metadata['FISH']:
+        count = 0
+        for frame in result['frames'][fish['START_FRAME']:fish['END_FRAME']+1]:
+            for ann in frame['fish']:
+                if ann['fish_id'] == fish['TOTAL']:
+                    count += 1
+        fish['DETECTION_DROPOUT'] = 1 - count / (fish['END_FRAME'] + 1 - fish['START_FRAME'])
+
     # Create fish table
     table = []
     for fish in metadata["FISH"]:

gradio_scripts/pdf_handler.py

@@ -0,0 +1,312 @@
+import datetime
+import numpy as np
+from matplotlib.backends.backend_pdf import PdfPages
+from matplotlib import collections as mc
+import matplotlib.pyplot as plt
+import math
+from aris import BEAM_WIDTH_DIR
+import cv2
+
+from dataloader import create_dataloader_aris
+
+
+STANDARD_FIG_SIZE = (16, 9)
+OUT_PDF_FILE_NAME = 'multipage_pdf.pdf'
+
+
+def make_pdf(i, state, result, table_headers):
+
+    fish_info = result["fish_info"][i]
+    fish_table = result["fish_table"][i]
+    json_result = result['json_result'][i]
+    metadata = json_result['metadata']
+    aris_input = result["aris_input"][i]
+
+    with PdfPages(OUT_PDF_FILE_NAME) as pdf:
+        plt.rcParams['text.usetex'] = False
+        
+        generate_title_page(pdf, metadata, state)
+        
+        generate_global_result(pdf, fish_info)
+
+        generate_fish_list(pdf, table_headers, fish_table)
+
+
+        dataset = None
+        if (aris_input is not None):
+            dataloader, dataset = create_dataloader_aris(aris_input, BEAM_WIDTH_DIR, None)
+
+        for i, fish in enumerate(json_result['fish']):
+            calculate_fish_paths(json_result, dataset, i)
+        
+        draw_combined_fish_graphs(pdf, json_result)
+
+        for i, fish in enumerate(json_result['fish']):
+            generate_fish_tracks(pdf, json_result, i)
+
+        # We can also set the file's metadata via the PdfPages object:
+        d = pdf.infodict()
+        d['Title'] = 'Multipage PDF Example'
+        d['Author'] = 'Oskar Åström'
+        d['Subject'] = 'How to create a multipage pdf file and set its metadata'
+        d['Keywords'] = ''
+        d['CreationDate'] = datetime.datetime.today()
+        d['ModDate'] = datetime.datetime.today()
+
+
+def generate_title_page(pdf, metadata, state):
+    # set up figure that will be used to display the opening banner
+    fig = plt.figure(figsize=STANDARD_FIG_SIZE)
+    plt.axis('off')
+
+    title_font_size = 40
+    minor_font_size = 20
+
+    # stuff to be printed out on the first page of the report
+    plt.text(0.5,-0.5,f'{metadata["FILE_NAME"].split("/")[-1]}',fontsize=title_font_size, horizontalalignment='center')
+
+    plt.text(0,1,f'Duration: {metadata["TOTAL_TIME"]}',fontsize=minor_font_size)
+    plt.text(0,1.5,f'Frames: {metadata["TOTAL_FRAMES"]}',fontsize=minor_font_size)
+    plt.text(0,2,f'Frame Rate: {metadata["FRAME_RATE"]}',fontsize=minor_font_size)
+
+    plt.text(0.5,1,f'Time of filming: {metadata["DATE"]} ({metadata["START"]} - {metadata["END"]})',fontsize=minor_font_size)
+    plt.text(0.5,1.5,f'Web app version: {state["version"]}',fontsize=minor_font_size)
+
+    plt.text(1.1,4.5,f'PDF generated on {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',fontsize=minor_font_size, horizontalalignment='right')
+
+    plt.ylim([-1, 4])
+    plt.xlim([0, 1])
+    plt.gca().invert_yaxis()
+    
+    pdf.savefig(fig)
+    plt.close(fig)
+
+def generate_global_result(pdf, fish_info):
+    # set up figure that will be used to display the opening banner
+    fig = plt.figure(figsize=STANDARD_FIG_SIZE)
+    plt.axis('off')
+    # stuff to be printed out on the first page of the report
+
+    minor_font_size = 18
+
+    headers = ["Result", "Camera Info", "Hyperparameters"]
+    info_col_1 = []
+    info_col_2 = []
+    info_col = info_col_1
+    row_state = -1
+    for row in fish_info:
+        if row_state >= 0:
+            info_col.append([row[0].replace("**","").replace("_", " ").lower(), row[1], 'normal'])
+        if (row[0] == "****"):
+            row_state += 1
+            if row_state == 2: info_col = info_col_2
+            info_col.append([headers[row_state], "", 'bold'])
+    for row_i, row in enumerate(info_col_1):
+        h = -1 + 5*row_i/len(info_col_1)
+        plt.text(0,    h, row[0], fontsize=minor_font_size, weight=row[2])
+        plt.text(0.25, h, row[1], fontsize=minor_font_size, weight=row[2])
+    for row_i, row in enumerate(info_col_2):
+        h = -1 + 5*row_i/len(info_col_2)
+        plt.text(0.5,  h, row[0], fontsize=minor_font_size, weight=row[2])
+        plt.text(0.75, h, row[1], fontsize=minor_font_size, weight=row[2])
+    plt.ylim([-1, 4])
+    plt.xlim([0, 1])
+    plt.gca().invert_yaxis()
+    
+    pdf.savefig(fig)
+    plt.close(fig)
+
+def generate_fish_list(pdf, table_headers, fish_table):
+    # set up figure that will be used to display the opening banner
+    fig = plt.figure(figsize=STANDARD_FIG_SIZE)
+    plt.axis('off')
+    # stuff to be printed out on the first page of the report
+
+    title_font_size = 40
+    header_font_size = 12
+    body_font_size = 20
+
+    # Title
+    plt.text(0.5,-1.3,f'{"Identified Fish"}',fontsize=title_font_size, horizontalalignment='center', weight='bold')
+
+    # Identified fish
+    row_h = 0.25
+    col_start = 0
+    row_l = 1
+    dropout_i = None
+    for col_i, col in enumerate(table_headers):
+        x = col_start + row_l*(col_i+0.5)/len(table_headers)
+        if col == "TOTAL": col = "ID"
+        if col == "DETECTION_DROPOUT": 
+            col = "frame drop rate"
+            dropout_i = col_i
+        col = col.lower().replace("_", " ")
+        plt.text(x, -1, col, fontsize=header_font_size, horizontalalignment='center', weight="bold")
+        plt.plot([col_start*2, -col_start*2 + row_l], [-1 + 0.05, -1 + 0.05], color='black')
+
+    for row_i, row in enumerate(fish_table):
+        y = -1 + (row_i+1)*row_h
+        plt.plot([col_start*2, -col_start*2 + row_l], [y+0.05, y+0.05], color='black')
+        for col_i, col in enumerate(row):
+            x = col_start + row_l*(col_i+0.5)/len(row)
+            if (col_i == dropout_i):
+                col = str(int(col*100)) + "%"
+            elif type(col) == float:
+                col = "{:.4f}".format(col)
+            plt.text(x, y, col, fontsize=body_font_size, horizontalalignment='center')
+    plt.ylim([-1, 4])
+    plt.xlim([0, 1])
+    plt.gca().invert_yaxis()
+    
+    pdf.savefig(fig)
+    plt.close(fig)
+
+def calculate_fish_paths(result, dataset, id):
+
+    fish = result['metadata']['FISH'][id]
+    start_frame = fish['START_FRAME']
+    end_frame = fish['END_FRAME']
+
+    # Extract base frame (first frame for that fish)
+    w, h = 1, 2
+    img = None
+    if (dataset is not None):
+        
+        images = dataset.didson.load_frames(start_frame=start_frame, end_frame=start_frame+1)
+        img = images[0]
+
+        frame_height = 2
+        scale_factor = frame_height / h
+        h = frame_height
+        w = int(scale_factor*w)
+
+    fish['base_frame'] = img
+    fish['scaled_frame_size'] = (w, h)
+
+
+    # Find frames for this fish
+    bboxes = []
+    for frame in result['frames'][start_frame:end_frame+1]:
+        bbox = None
+        for ann in frame['fish']:
+            if ann['fish_id'] == id+1:
+                bbox = ann
+        bboxes.append(bbox)
+
+
+    # Calculate tracks through frames
+    missed = 0
+    X = []
+    Y = []
+    V = []
+    certainty = []
+    for bbox in bboxes:
+        if bbox is not None:
+
+            # Find fish centers
+            x = (bbox['bbox'][0] + bbox['bbox'][2])/2
+            y = (bbox['bbox'][1] + bbox['bbox'][3])/2
+
+            # Calculate velocity
+            v = None
+            if len(X) > 0:
+                last_x = X[-1]
+                last_y = Y[-1]
+                dx = result['image_meter_width']*(last_x - x)/(missed+1)
+                dy = result['image_meter_height']*(last_y - y)/(missed+1)
+                v = math.sqrt(dx*dx + dy*dy)
+
+            # Interpolate over missing frames
+            if missed > 0:    
+                for i in range(missed):
+                    p = (i+1)/(missed+1)
+                    X.append(last_x*(1-p) + p*x)
+                    Y.append(last_y*(1-p) + p*y)
+                    V.append(v)
+                    certainty.append(False)
+            
+            # Append new track frame
+            X.append(x)
+            Y.append(y)
+            if v is not None: V.append(v)
+            certainty.append(True)
+            missed = 0
+        else:
+            missed += 1
+
+    fish['path'] = {
+        'X': X,
+        'Y': Y,
+        'certainty': certainty,
+        'V': V
+    }
+
+
+def draw_combined_fish_graphs(pdf, result):
+    vel = []
+    log_vel = []
+    for fish in result['metadata']['FISH']:
+        vel += fish['path']['V']
+        log_vel += [math.log(v) for v in fish['path']['V']]
+
+    fig, axs = plt.subplots(2, 2, sharey=True, figsize=STANDARD_FIG_SIZE)
+    axs[0,0].hist(log_vel, bins=20)
+    axs[0,0].set_title('Fish Log-Velocities between frames')
+    axs[0,0].set_xlabel("Log-Velocity (log(m/frame))")
+    axs[0,1].hist(vel, bins=20)
+    axs[0,1].set_title('Fish Velocities between frames')
+    axs[0,1].set_xlabel("Velocity (m/frame)")
+
+    pdf.savefig(fig)
+    plt.close(fig)
+    
+
+def generate_fish_tracks(pdf, result, id):
+
+    fish = result['metadata']['FISH'][id]
+    start_frame = fish['START_FRAME']
+    end_frame = fish['END_FRAME']
+
+    fig, ax = plt.subplots(figsize=STANDARD_FIG_SIZE)
+    plt.axis('off')
+
+    w, h = fish['scaled_frame_size']
+    if (fish['base_frame'] is not None):
+        img = fish['base_frame']
+        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
+        plt.imshow(img, extent=(0, h, 0, w), cmap=plt.colormaps['Greys'])
+
+    # Title
+    plt.text(h/2,1.1,f'Fish {id+1} (frames {start_frame}-{end_frame})',fontsize=40, color="red", horizontalalignment='center', zorder=5)
+
+    X = fish['path']['X']
+    Y = fish['path']['Y']
+    certainty = fish['path']['certainty']
+    
+    plt.text(h*(1-Y[0]), w*(1-X[0]), "Start", fontsize=15, weight="bold")
+    plt.text(h*(1-Y[-1]), w*(1-X[-1]), "End", fontsize=15, weight="bold")
+
+    colors = []
+    for i in range(1, len(X)):
+
+        certain = certainty[i]
+        fully_certain = certain
+        half_certain = certain
+        if i > 0: 
+            fully_certain &= certainty[i-1]
+            half_certain |= certainty[i-1]
+            
+        #color = 'yellow' if certain else 'orangered'
+        #plt.plot(h*(1-y), w*(1-x), marker='o', markersize=3, color=color, zorder=3)
+        col = 'yellow' if fully_certain else ('darkorange' if half_certain else 'orangered')
+        colors.append(col)
+        ax.plot([h*(1-Y[i-1]), h*(1-Y[i])], [w*(1-X[i-1]), w*(1-X[i])], color=col, linewidth=1)
+    
+    for i in range(1, len(X)):
+        ax.plot(h*(1-Y[i]), w*(1-X[i]), color=colors[i], marker='o', markersize=3)
+    
+
+    plt.ylim([0, w])
+    plt.xlim([0, h])
+    pdf.savefig(fig)
+    plt.close(fig)

gradio_scripts/result_ui.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import numpy as np
+from gradio_scripts.pdf_handler import make_pdf
 
 js_update_tab_labels = """
     async () => {
@@ -14,7 +15,7 @@ js_update_tab_labels = """
     }
 """
 
-table_headers = ["TOTAL", "START_FRAME", "END_FRAME", "DIR", "R", "THETA", "L", "TRAVEL"]
+table_headers = ["TOTAL", "START_FRAME", "END_FRAME", "DETECTION_DROPOUT", "DIR", "R", "THETA", "L", "TRAVEL"]
 info_headers = [
     "TOTAL_TIME", "DATE", "START", "END", "FRAME_RATE", "",
     "TOTAL_FISH", "UPSTREAM_FISH", "DOWNSTREAM_FISH", "NONDIRECTIONAL_FISH", "",
@@ -38,6 +39,8 @@ def update_result(i, state, result, inference_handler):
 
     annotation_avaliable = not (result["aris_input"][i] == None)
 
+    make_pdf(state['index']-1, state, result, table_headers)
+
     # Send update to UI, and to inference_handler to start next file inference
     return {
         zip_out: gr.update(value=result["path_zip"]),

gradio_scripts/upload_ui.py

@@ -31,12 +31,16 @@ def Upload_Gradio(gradio_components):
                     settings.append(gr.Slider(0, 100, value=16, label="Min Hits", info="Minimum number of frames a fish has to appear in to count"))
                     settings.append(gr.Slider(0, 100, value=14, label="Max Age", info="Max age of occlusion before track is split"))
 
-                with gr.Row():
-                    gr.Markdown("Associative Tracking")
-                    settings.append(gr.Checkbox(value=False, label="Enabled"))
-                with gr.Row():
+                tracker = gr.Dropdown(["None", "Confidence Boost", "ByteTrack"], label="Associative Tracking", value="None")
+                settings.append(tracker)
+                with gr.Row(visible=False) as track_row:
                     settings.append(gr.Slider(0, 5, value=1, label="Boost Power", info=""))
                     settings.append(gr.Slider(0, 1, value=1, label="Boost Decay", info=""))
+                    tracker.change(lambda x:  gr.update(visible=(x=="Confidence Boost")), tracker, track_row)
+                with gr.Row(visible=False) as track_row:
+                    settings.append(gr.Slider(0, 1, value=0.1, label="Low Conf Threshold", info=""))
+                    settings.append(gr.Slider(0, 1, value=0.3, label="High Conf Threshold", info=""))
+                    tracker.change(lambda x:  gr.update(visible=(x=="ByteTrack")), tracker, track_row)
 
                 gr.Markdown("Other")
                 with gr.Row():

inference.py

@@ -58,12 +58,9 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
     if 'iou_thresh' not in hyperparams: hyperparams['iou_thresh'] = NMS_IOU
     if 'min_hits' not in hyperparams: hyperparams['min_hits'] = MIN_HITS
     if 'max_age' not in hyperparams: hyperparams['max_age'] = MAX_AGE
-    if 'use_associative_tracking' not in hyperparams: hyperparams['use_associative_tracking'] = False
-    if 'boost_power' not in hyperparams: hyperparams['boost_power'] = 1
-    if 'boost_decay' not in hyperparams: hyperparams['maxboost_decay_age'] = 1
-    if 'AT_decay' not in hyperparams: hyperparams['AT_decay'] = MIN_HITS
     if 'min_length' not in hyperparams: hyperparams['min_length'] = MIN_LENGTH
     if 'min_travel' not in hyperparams: hyperparams['min_travel'] = MIN_TRAVEL
+    if 'associative_tracker' not in hyperparams: hyperparams['associative_tracker'] = "None"
     
     model, device = setup_model(hyperparams['model'])
 
@@ -95,16 +92,41 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
 
     outputs = do_suppression(inference, conf_thres=hyperparams['conf_thresh'], iou_thres=hyperparams['iou_thresh'], gp=gp)
 
-    if hyperparams['use_associative_tracking']:
+    if hyperparams['associative_tracker'] == "ByteTrack":
+        if 'byte_low_conf' not in hyperparams: hyperparams['byte_low_conf'] = 0.1
+        if 'byte_high_conf' not in hyperparams: hyperparams['byte_high_conf'] = 0.3
 
-        do_confidence_boost(inference, outputs, conf_power=hyperparams['boost_power'], conf_decay=hyperparams['boost_decay'], gp=gp)
+        low_outputs = do_suppression(inference, conf_thres=hyperparams['low_conf_threshold'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+        low_preds, real_width, real_height = format_predictions(image_shapes, low_outputs, width, height, gp=gp)
 
-        outputs = do_suppression(inference, conf_thres=hyperparams['conf_thresh'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+        high_outputs = do_suppression(inference, conf_thres=hyperparams['high_conf_threshold'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+        high_preds, real_width, real_height = format_predictions(image_shapes, high_outputs, width, height, gp=gp)
 
-    all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, gp=gp)
+        results = do_associative_tracking(
+            low_preds, high_preds, image_meter_width, image_meter_height, 
+            reverse=False, min_length=hyperparams['min_length'], min_travel=hyperparams['min_travel'], 
+            max_age=hyperparams['max_age'], min_hits=hyperparams['min_hits'], 
+            gp=gp)
+    else: 
+
+        outputs = do_suppression(inference, conf_thres=hyperparams['conf_threshold'], iou_thres=hyperparams['iou_thresh'], gp=gp)
+
+        if hyperparams['associative_tracker'] == "Confidence Boost":
+            if 'boost_power' not in hyperparams: hyperparams['boost_power'] = 1
+            if 'boost_decay' not in hyperparams: hyperparams['boost_decay'] = 1
+
+            do_confidence_boost(inference, outputs, boost_power=hyperparams['boost_power'], boost_decay=hyperparams['boost_decay'], gp=gp)
+
+            outputs = do_suppression(inference, conf_thres=hyperparams['conf_threshold'], iou_thres=hyperparams['iou_thresh'], gp=gp)
 
-    results = do_tracking(all_preds, image_meter_width, image_meter_height, min_hits=hyperparams['min_hits'], max_age=hyperparams['max_age'], min_length=hyperparams['min_length'], min_travel=hyperparams['min_travel'], gp=gp)
+        all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height, gp=gp)
 
+        results = do_tracking(
+            all_preds, image_meter_width, image_meter_height, 
+            min_length=hyperparams['min_length'], min_travel=hyperparams['min_travel'], 
+            max_age=hyperparams['max_age'], iou_thres=hyperparams['iou_threshold'], min_hits=hyperparams['min_hits'], 
+            gp=gp)
+ 
     return results