First Commit

.gitignore

@@ -0,0 +1,18 @@
+venv/
+.vscode
+#*.pt
+#*.jpg
+static/tmp.jpg
+redis-stable/*
+user_data/*
+*.pyc
+
+.ipynb_checkpoints
+.tmp*
+*.mp4
+*.jpg
+*.json
+*.zip
+*.aris
+*.log
+*.DS_STORE

app.py

@@ -0,0 +1,212 @@
+import gradio as gr
+from uploader import save_data
+from main import predict_task
+from state_handler import load_example_result, reset_state
+from file_reader import File
+import numpy as np
+from aws_handler import upload_file
+from aris import create_metadata_table
+
+table_headers = ["TOTAL" , "FRAME_NUM", "DIR", "R", "THETA", "L", "TIME", "DATE", "SPECIES"]
+info_headers = [
+    "TOTAL_TIME", "DATE", "START", "END",
+    "TOTAL_FISH", "UPSTREAM_FISH", "DOWNSTREAM_FISH", "NONDIRECTIONAL_FISH", 
+    "TOTAL_FRAMES", "FRAME_RATE", 
+    "UPSTREAM_MOTION", "INTENSITY", "THRESHOLD", "WINDOW_START", "WINDOW_END", "WATER_TEMP"
+]
+css = """
+#result_json {
+    height: 500px;
+    overflow: scroll !important;
+}
+#marking_json textarea {
+    height: 100% !important;
+}
+#marking_json label {
+    height: calc(100% - 30px) !important;
+}
+"""
+js_update_tabs = """
+    async () => {
+        let el_list = document.getElementById("result_handler").getElementsByClassName("svelte-1kcgrqr")
+        let idx = (el_list[1].value === "LOADING") ? 1 : parseInt(el_list[1].value)
+        console.log(idx)
+        style_sheet = document.getElementById("tab_style")
+        style_sheet.innerHTML = ""
+        for (let i = 1; i <= idx; i++) {
+            style_sheet.innerHTML += "button.svelte-kqij2n:nth-child(" + i + "):before {content: 'Result " + i + "';}"
+        }
+    }
+"""
+
+#Initialize State & Result
+state = {
+    'files': [],
+    'index': 1,
+    'total': 1
+}
+result = {}
+
+# Start function, called on file upload
+def on_input(file_list):
+
+    # Reset Result
+    reset_state(result, state)
+    state['files'] = file_list
+    state['total'] = len(file_list)
+
+    # Update loading_space to start inference on first file
+    return {
+        inference_handler: gr.update(value = str(np.random.rand()), visible=True)
+    }
+
+# Iterative function that performs inference on the next file in line
+def handle_next(_, progress=gr.Progress()):
+
+    if state['index'] >= state['total']: 
+        return {
+            result_handler: gr.update(),
+            inference_handler: gr.update()
+        }
+
+    # Correct progress function for batch file input
+    set_progress = lambda pct, msg : progress(pct, desc=msg)
+    if state['total'] > 1:
+        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]
+    valid, file_path, dir_name = save_data(bytes, file_name)
+    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)
+    result["path_marking"].append(marking_filepath)
+    fish_table, fish_info = create_metadata_table(metadata, table_headers, info_headers)
+    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()
+    }
+
+# Show result UI based on example data
+def show_example_data():
+    load_example_result(result, table_headers, info_headers)
+    state["index"] = 1
+    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']
+    message = "Result " + str(state['index']) + "/" + str(state['total'])
+    return {
+        zip_out: gr.update(value=result["path_zip"]),
+        tabs[i]['tab']: gr.update(),
+        tabs[i]['video']: gr.update(value=result["path_video"][i], visible=True),
+        tabs[i]['metadata']: gr.update(value=result["fish_info"][i], visible=True),
+        tabs[i]['table']: gr.update(value=result["fish_table"][i], visible=True),
+        tab_parent: gr.update(selected=i),
+        inference_handler: gr.update(value = str(np.random.rand()), visible=not_done)
+    }
+
+
+
+max_tabs = 10
+demo = gr.Blocks()
+with demo:
+    with gr.Blocks(css=css) as inner_body:
+
+        # Title of page
+        gr.HTML(
+            """
+        <h1 align="center" style="font-size:xxx-large">Caltech Fisheye</h1>
+        <p align="center">Submit an .aris file to analyze result.</p>
+        <style id="tab_style"></style>
+        """
+        )
+        
+        #Input field for aris submission
+        input = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")
+
+        # Dummy element to call inference events, this also displays the inference progress
+        inference_handler = gr.Text(value=str(np.random.rand()), visible=False)
+
+        # Dummy element to call UI events
+        result_handler = gr.Text(value="LOADING", visible=False, elem_id="result_handler")
+
+        # List of all UI components that will recieve outputs from the result_handler
+        UI_components = []
+
+        # Zip file output
+        zip_out = gr.File(label="ZIP Output", interactive=False)
+        UI_components.append(zip_out)
+
+        # Create result tabs
+        tabs = []
+        with gr.Tabs() as tab_parent:
+            UI_components.append(tab_parent)
+            for i in range(max_tabs):
+                with gr.Tab(label="", id=i, elem_id="result_tab"+str(i)) as tab:
+                    with gr.Row():
+                        metadata_out = gr.JSON(label="Info", visible=False, elem_id="marking_json")
+                        video_out = gr.Video(label='Annotated Video', interactive=False, visible=False)
+                    table_out = gr.Matrix(label='Indentified Fish', headers=table_headers, interactive=False, visible=False)
+                    tabs.append({
+                        'tab': tab,
+                        'metadata': metadata_out,
+                        'video': video_out,
+                        'table': table_out
+                    })
+                    UI_components.extend([tab, metadata_out, video_out, table_out])
+
+    # Button to show example result
+    gr.Button(value="Show Example Result").click(show_example_data, None, result_handler)
+
+    # Disclaimer at the bottom of page
+    gr.HTML(
+            """
+        <p align="center">
+        <b>Note</b>: The software is provided "as is", without warranty of any kind, express or implied, including but not limited to the warranties of merchantability, fitness for a particular purpose and noninfringement. 
+        In no event shall the authors or copyright holders be liable for any claim, damages or other liability, whether in an action of contract, tort or otherwise, arising from, out of or in connection with the software or the use or other dealings in the software.
+        </p>
+        """
+    )
+
+    # When a file is uploaded to the input, tell the inference_handler to start inference
+    input.upload(fn=on_input, inputs=input, outputs=[inference_handler])
+
+    # When inference handler updates, tell result_handler to show the new result
+    # Also, add inference_handler as the output in order to have it display the progress
+    inference_handler.change(handle_next, None, [result_handler, inference_handler])
+
+    # Send UI changes based on the new results to the UI_components, and tell the inference_handler to start next inference
+    result_handler.change(show_data, None, UI_components + [inference_handler], _js=js_update_tabs)
+
+demo.queue().launch()
+
+show_data()
+
+

aris.py

@@ -0,0 +1,505 @@
+import project_path
+
+import numpy as np
+import cv2
+import os
+from collections import namedtuple, defaultdict
+import struct
+from PIL import Image
+from tqdm import tqdm
+import datetime
+from decimal import Decimal, ROUND_HALF_UP
+import json
+import pytz
+from copy import deepcopy
+from multiprocessing import Pool
+
+import pyARIS
+from tracker import Tracker
+
+
+BEAM_WIDTH_DIR = 'lib/fish_eye/beam_widths/'
+
+ImageData = namedtuple('ImageData', [
+    'pixel_meter_size',
+    'xdim', 'ydim',
+    'x_meter_start', 'y_meter_start', 'x_meter_stop', 'y_meter_stop',
+    'sample_read_rows', 'sample_read_cols', 'image_write_rows', 'image_write_cols'
+])
+
+def FastARISRead(ARIS_data, start_frame, end_frame):
+    """ Just read in the ARIS frame, and not the other meta data. 
+    """
+    FrameSize = ARIS_data.SamplesPerChannel*ARIS_data.NumRawBeams
+    frames = np.empty([end_frame-start_frame, ARIS_data.SamplesPerChannel,
+                      ARIS_data.NumRawBeams], dtype=np.uint8)
+    with open(ARIS_data.filename, 'rb') as data:
+        for i, j in enumerate(range(start_frame, end_frame)):
+            data.seek(j*(1024+(FrameSize))+2048, 0)
+            raw_data = struct.unpack("%dB" % FrameSize, data.read(FrameSize))
+            frames[i] = np.fliplr(np.reshape(
+                raw_data, [ARIS_data.SamplesPerChannel, ARIS_data.NumRawBeams]))
+    # Close the data file
+    data.close()
+    return frames
+
+def get_info(aris_fp, beam_width_dir=BEAM_WIDTH_DIR):
+    """
+    Return:
+        image_meter_width, image_meter_height, fps
+    """
+    ARISdata, aris_frame = pyARIS.DataImport(aris_fp)
+    beam_width_data = pyARIS.load_beam_width_data(aris_frame, beam_width_dir=beam_width_dir)[0]
+    min_pixel_size = pyARIS.get_minimum_pixel_meter_size(aris_frame, beam_width_data)
+    sample_length = aris_frame.sampleperiod * 0.000001 * aris_frame.soundspeed / 2
+    pixel_meter_size = max(min_pixel_size, sample_length)
+    xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds(
+        pixel_meter_size, aris_frame, beam_width_data,
+        additional_pixel_padding_x=0,
+        additional_pixel_padding_y=0
+    )
+    return pixel_meter_size * xdim, pixel_meter_size * ydim, aris_frame.framerate
+
+def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH_DIR, bg_out_dir=None, num_workers=0):
+    """
+    Write all frames from an ARIS file to disk, using our 3-channel format:
+    (raw img, blurred & mean subtracted img, optical flow approximation)
+    
+    Args:
+        aris_fp: path to aris file
+        out_dir: directory for frame extraction. frames will be named 0.jpg, 1.jpg, ... {n}.jpg
+        cb: a callback function for updating progress
+        max_mb: maximum amount of the file to be processed, in megabytes
+        beam_width_dir: location of ARIS camera information
+        bg_out_dir: where to write the background frame; None disables writing
+    
+    Return:
+        (float) image_meter_width - the width of each image, in meters
+        (float) image_meter_height
+        (float) fps
+    """
+    # Load in the ARIS file
+    ARISdata, aris_frame = pyARIS.DataImport(aris_fp)
+    
+    if cb:
+        cb(2, msg="Decoding ARIS data...")
+    
+    beam_width_data = pyARIS.load_beam_width_data(aris_frame, beam_width_dir=beam_width_dir)[0]
+    # What is the meter resolution of the smallest sample?
+    min_pixel_size = pyARIS.get_minimum_pixel_meter_size(aris_frame, beam_width_data)
+    # What is the meter resolution of the sample length?
+    sample_length = aris_frame.sampleperiod * 0.000001 * aris_frame.soundspeed / 2
+    # Choose the size of a pixel (or hard code it to some specific value)
+    pixel_meter_size = max(min_pixel_size, sample_length)
+    # Determine the image dimensions
+    xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds(
+        pixel_meter_size, aris_frame, beam_width_data,
+        additional_pixel_padding_x=0,
+        additional_pixel_padding_y=0
+    )
+    
+    # Compute the mapping from the samples to the image
+    sample_read_rows, sample_read_cols, image_write_rows, image_write_cols = pyARIS.compute_mapping_from_sample_to_image(
+        pixel_meter_size,
+        xdim, ydim, x_meter_start, y_meter_start,
+        aris_frame, beam_width_data
+    )
+    image_data = ImageData(
+        pixel_meter_size,
+        xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop,
+        sample_read_rows, sample_read_cols, image_write_rows, image_write_cols
+    )
+    start_frame = 0
+    end_frame = ARISdata.FrameCount
+    
+    bytes_per_frame = 1024+ARISdata.SamplesPerChannel*ARISdata.NumRawBeams
+    print("ARIS bytes per frame", bytes_per_frame)
+    
+    img_bytes_per_frame = image_data.ydim * image_data.xdim * 4 # for fp32 frames
+    print("Image bytes per frame", img_bytes_per_frame)
+      
+    max_bytes = max(bytes_per_frame, img_bytes_per_frame)
+        
+    if max_mb > 0:
+        max_frames = int(max_mb*1000000 / (max_bytes))
+        if end_frame > max_frames:
+            end_frame = max_frames
+        
+    # use a max of 4gb per batch to avoid memory errors (16gb RAM on a g4dn.xlarge)
+    batch_size = 1000 # int(4000*1000000 / (max_bytes))
+    clips = [[pos, pos+batch_size+1] for pos in range(0, end_frame, batch_size)]
+    clips[-1][1] = ARISdata.FrameCount
+    print("Batch size:", batch_size)
+    
+    with tqdm(total=(end_frame-start_frame-1), desc="Extracting frames", ncols=0) as pbar:
+        # compute info for bg subtraction using first batch
+        # TODO: make this a sliding window
+        mean_blurred_frame, mean_normalization_value = write_frame_range(ARISdata, image_data, out_dir, clips[0][0], clips[0][1], None, None, cb, pbar)
+    
+        # do rest of batches in parallel
+        if num_workers > 0:
+            args = [ (ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, cb) for (start, end) in clips[1:] ] # TODO: can't pass pbar to thread
+            with Pool(num_workers) as pool:
+                results = [ pool.apply_async(write_frame_range, arg) for arg in args ]
+                results = [ r.get() for r in results ] # need this call to block on thread execution
+                pbar.update(sum([ arg[4] - arg[3] for arg in args ]))
+        else:
+            for j, (start, end) in enumerate(clips[1:]):
+                write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, cb, pbar)
+            
+    if bg_out_dir is not None:
+        bg_img = (mean_blurred_frame * 255).astype(np.uint8)
+        out_fp = os.path.join(bg_out_dir, 'bg_start.jpg')
+        Image.fromarray(bg_img).save(out_fp, quality=95)
+        
+    return pixel_meter_size * xdim, pixel_meter_size * ydim, aris_frame.framerate
+
+def write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame=None, mean_normalization_value=None, cb=None, pbar=None):
+    try:
+        frames = np.zeros([end-start, image_data.ydim, image_data.xdim], dtype=np.uint8)
+        frames[:, image_data.image_write_rows, image_data.image_write_cols] = FastARISRead(ARISdata, start, end)[:, image_data.sample_read_rows, image_data.sample_read_cols]
+    except:
+        print("Error extracting frames from", ARISdata.filename, "during batch", i)
+        return
+
+    blurred_frames = frames.astype(np.float32)
+    for i in range(frames.shape[0]):
+        blurred_frames[i] = cv2.GaussianBlur(
+            blurred_frames[i],
+            (5,5),
+            0
+        )
+
+    if mean_blurred_frame is None:
+        mean_blurred_frame = blurred_frames.mean(axis=0)
+
+    blurred_frames -= mean_blurred_frame
+
+    if mean_normalization_value is None:
+        mean_normalization_value = np.max(np.abs(blurred_frames))
+
+    blurred_frames /= mean_normalization_value
+    blurred_frames += 1
+    blurred_frames /= 2
+
+    # Because of the optical flow computation, we only go to end_frame - 1
+    for i, frame_offset in enumerate(range(start, end - 1)):
+        frame_image = np.dstack([
+                        frames[i] / 255,
+                        blurred_frames[i],
+                        np.abs(blurred_frames[i+1] - blurred_frames[i])
+                    ]).astype(np.float32)
+        frame_image = (frame_image * 255).astype(np.uint8)
+        out_fp = os.path.join(out_dir, f'{start+i}.jpg') # = frame_offset.jpg?
+        Image.fromarray(frame_image).save(out_fp, quality=95)
+        
+        if pbar:
+            pbar.update(1)
+        if cb:
+            pct = 2 + int( (start+i) / (end_frame - start_frame - 1) * 98)
+            cb(pct, msg=pbar.__str__())
+            
+    return mean_blurred_frame, mean_normalization_value
+
+def prep_for_mm(json_data):
+    """Prepare json results for writing to a manual marking file."""
+    json_data = deepcopy(json_data)
+            
+    # map fish id -> [ (bbox, frame_num), (bbox, frame_num), ... ]
+    tracks = defaultdict(list)
+    for frame in json_data['frames']:
+        for bbox in frame['fish']:
+            tracks[bbox['fish_id']].append((bbox['bbox'], frame['frame_num']))
+    
+    # find frame number for manual marking
+    # look for first time a track crosses the center
+    # if it never crosses the center, use the closest box to the center
+    mm_frame_nums = {}
+    for f_id, track in tracks.items():
+        # keep track of frame closest to the center
+        closest_frame = 0 
+        closest_dist = 1.0
+        for i, (box, frame) in enumerate(track):
+            x = (box[0] + box[2]) / 2.0
+            if i > 0:
+                last_x = (track[i-1][0][0] + track[i-1][0][2]) / 2.0
+                if (x < 0.5 and last_x >= 0.5) or (last_x < 0.5 and x >= 0.5):
+                    closest_frame = frame
+                    break
+            dist = abs(x - 0.5)
+            if dist < closest_dist:
+                closest_frame = frame
+                closest_dist = dist
+        mm_frame_nums[f_id] = closest_frame
+        
+    # sort tracks by their frame numbers and re-key
+    # IDs are 1-indexed
+    id_frame = [ (k, v) for k,v in mm_frame_nums.items() ]
+    id_frame = sorted(id_frame, key=lambda x: x[1])
+    id_map = {}
+    for i, (f_id, frame) in enumerate(id_frame, start=1):
+        id_map[f_id] = i
+        
+    # map IDs and keep frame['fish'] sorted by ID
+    for i, frame in enumerate(json_data['frames']):
+        new_frame_entries = []
+        for frame_entry in frame['fish']:
+            frame_entry['fish_id'] = id_map[frame_entry['fish_id']]
+            new_frame_entries.append(frame_entry)
+        frame['fish'] = sorted(new_frame_entries, key=lambda k: k['fish_id'])
+    
+    # store manual marking frame and re-map 'fish' field
+    for fish in json_data['fish']:
+        fish['marking_frame'] = mm_frame_nums[fish['id']] # mm_frame_nums refers to old IDs
+        fish['id'] = id_map[fish['id']]
+    json_data['fish'] = sorted(json_data['fish'], key=lambda x: x['id'])
+        
+    return json_data
+
+
+def create_metadata_dictionary(aris_fp, json_fp, beam_width_dir=BEAM_WIDTH_DIR):
+    """
+    Return:
+        dictionary, for manual marking
+    """
+    with open(json_fp) as json_file:
+        json_data = json.load(json_file)
+
+    metadata = {}
+    metadata["FILE_NAME"] = aris_fp
+    ARISdata, frame = pyARIS.DataImport(aris_fp)
+    metadata["FRAME_RATE"] = frame.framerate
+
+    # Load in the beam width information
+    beam_width_data, camera_type = pyARIS.load_beam_width_data(frame, beam_width_dir=beam_width_dir)
+
+    # What is the meter resolution of the smallest sample?
+    min_pixel_size = pyARIS.get_minimum_pixel_meter_size(frame, beam_width_data)
+
+    # What is the meter resolution of the sample length?
+    sample_length = frame.sampleperiod * 0.000001 * frame.soundspeed / 2
+
+    # Choose the size of a pixel
+    pixel_meter_size = max(min_pixel_size, sample_length)
+
+    # Determine the image dimensions
+    xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds(
+        pixel_meter_size, frame, beam_width_data,
+        additional_pixel_padding_x=0,
+        additional_pixel_padding_y=0
+    )
+
+    # Compute the mapping from the samples to the image
+    sample_read_rows, sample_read_cols, image_write_rows, image_write_cols = pyARIS.compute_mapping_from_sample_to_image(
+        pixel_meter_size,
+        xdim, ydim, x_meter_start, y_meter_start,
+        frame, beam_width_data
+    )
+
+    marking_mapping = dict(zip(zip(image_write_rows, image_write_cols),
+                               zip(sample_read_rows, sample_read_cols)))
+    
+    # Manual marking format rounds 0.5 to 1 instead of 0 in IEEE 754
+    def round(number, ndigits=0):
+        return float(Decimal(number).quantize(ndigits, ROUND_HALF_UP))
+
+    right, left, none = Tracker.count_dirs(json_data)
+    
+    metadata["UPSTREAM_FISH"] = left # TODO
+    metadata["DOWNSTREAM_FISH"] = right # TODO
+    metadata["NONDIRECTIONAL_FISH"] = none # TODO
+    metadata["TOTAL_FISH"] = metadata["UPSTREAM_FISH"] + metadata["DOWNSTREAM_FISH"] + metadata["NONDIRECTIONAL_FISH"]
+
+    metadata["TOTAL_FRAMES"] = ARISdata.FrameCount
+    metadata["EXPECTED_FRAMES"] = -1 # What is this?
+    metadata["TOTAL_TIME"] = str(datetime.timedelta(seconds=round(metadata["TOTAL_FRAMES"]/metadata["FRAME_RATE"])))
+    metadata["EXPECTED_TIME"] = str(datetime.timedelta(seconds=round(metadata["EXPECTED_FRAMES"]/metadata["FRAME_RATE"])))
+
+    metadata["UPSTREAM_MOTION"] = 'Right To Left' or 'Left To Right' #TODO
+
+    metadata["COUNT_FILE_NAME"] = 'N/A'
+    metadata["EDITOR_ID"] = 'N/A'
+    metadata["INTENSITY"] = f'{round(frame.threshold, 1):.1f} dB' # Missing
+    metadata["THRESHOLD"] = f'{round(frame.threshold, 1):.1f} dB' # Missing
+    metadata["WINDOW_START"] = round(frame.windowstart, 2)
+    metadata["WINDOW_END"] = round(frame.windowstart + frame.windowlength, 2)
+    metadata["WATER_TEMP"] = f'{int(round(frame.watertemp))} degC'
+
+    s = f''''''
+
+    upstream_motion_map = {}
+    if (metadata["UPSTREAM_MOTION"] == 'Left To Right'):
+        upstream_motion_map = {
+            'right': '  Up',
+            'left': 'Down',
+            'none': ' N/A',
+        }
+    elif (metadata["UPSTREAM_MOTION"] == 'Right To Left'):
+        upstream_motion_map = {
+            'left': '  Up',
+            'right': 'Down',
+            'none': ' N/A',
+        }
+
+    def get_entry(fish):
+        if 'marking_frame' in fish:
+            frame_num = fish['marking_frame']
+            entry = None
+            for json_frame in json_data['frames']:
+                if json_frame['frame_num'] == frame_num:
+                    for json_frame_entry in json_frame['fish']:
+                        if json_frame_entry['fish_id'] == fish['id']:
+                            json_frame_entry = json_frame_entry.copy()
+                            json_frame_entry['frame_num'] = frame_num
+                            return json_frame_entry
+        else:
+            print("Warning: JSON not correctly formatted for manual marking creation. Use aris.prep_for_mm()")
+            entries = []
+            for json_frame in json_data['frames']:
+                for json_frame_entry in json_frame['fish']:
+                    if json_frame_entry['fish_id'] == fish['id']:
+                        entries.append({'frame_num': json_frame['frame_num'], **json_frame_entry})
+            entry = entries[len(entries)//2]
+            return entry
+        print("Error, could not find entry for", fish)
+        return None # TODO better error handling
+    
+    entries = []
+    for fish in json_data['fish']:
+        entry = get_entry(fish)
+        entry['length'] = fish['length']*100
+        entry['direction'] = fish['direction']
+        entries.append(entry)
+    
+    metadata["FISH"] = []
+    for entry in sorted(entries, key=lambda x: x['fish_id']):
+        frame_num = entry['frame_num']
+        frame = pyARIS.FrameRead(ARISdata, frame_num)
+
+        y = (entry['bbox'][1]+entry['bbox'][3])/2
+        x = (entry['bbox'][0]+entry['bbox'][2])/2
+        h = np.max(image_write_rows)
+        w = np.max(image_write_cols)
+        # TODO actually fix this
+        try:
+            bin_num, beam_num = marking_mapping[(round(y*h), round(x*w))]
+        except:
+            bin_num = 0
+            beam_num = 0
+
+        fish_entry = {}
+        fish_entry['FILE'] = 1
+        fish_entry['TOTAL'] = entry['fish_id']
+        fish_entry['FRAME_NUM'] = entry['frame_num']
+        fish_entry['DIR'] = upstream_motion_map[entry['direction']]
+        fish_entry['R'] = bin_num * pixel_meter_size + frame.windowstart
+        fish_entry['THETA'] = beam_width_data['beam_center'][beam_num]
+        fish_entry['L'] = entry['length']
+        fish_entry['DR'] = -1.0 # What is this?
+        fish_entry['LDR'] = -1.0 # What is this?
+        fish_entry['ASPECT'] = -1.0 # What is this?
+        TIME, DATE = datetime.datetime.fromtimestamp(frame.sonartimestamp/1000000, pytz.timezone('UTC')).strftime('%H:%M:%S %Y-%m-%d').split()
+        fish_entry['TIME'] = TIME
+        fish_entry['DATE'] = DATE
+        fish_entry['LATITUDE'] = frame.latitude or 'N 00 d  0.00000 m'
+        fish_entry['LONGITUDE'] = frame.longitude or 'E 000 d  0.00000 m'
+        fish_entry['PAN'] = frame.sonarpan
+        fish_entry['TILT'] = frame.sonartilt
+        fish_entry['ROLL'] = frame.roll # May be wrong number but sonarroll was NaN
+        fish_entry['SPECIES'] = 'Unknown'
+        fish_entry['MOTION'] = 'Running <-->'
+        fish_entry['Q'] = -1 #5 # I don't know what this is or where it comes from
+        fish_entry['N'] = -1 #1 # I don't know what this is or where it comes from
+        fish_entry['COMMENT'] = ''
+        
+        metadata["FISH"].append(fish_entry)
+        
+    # What are these?
+    # Maybe the date and time range for the recording?
+    first_frame = pyARIS.FrameRead(ARISdata, 0)
+    last_frame = pyARIS.FrameRead(ARISdata, metadata["TOTAL_FRAMES"]-1)
+    start_time, start_date = datetime.datetime.fromtimestamp(first_frame.sonartimestamp/1000000, pytz.timezone('UTC')).strftime('%H:%M:%S %Y-%m-%d').split()
+    end_time, end_date = datetime.datetime.fromtimestamp(last_frame.sonartimestamp/1000000, pytz.timezone('UTC')).strftime('%H:%M:%S %Y-%m-%d').split()
+    metadata["DATE"] = start_date 
+    metadata["START"] = start_time
+    metadata["END"] = end_time
+
+    return metadata
+
+def create_metadata_table(metadata, table_headers, info_headers):
+    table = []
+    for fish in metadata["FISH"]:
+        row = []
+        for header in table_headers:
+            row.append(fish[header])
+        table.append(row)
+    
+    if len(metadata["FISH"]) == 0:
+        row = []
+        for header in table_headers:
+            row.append("-")
+        table.append(row)
+
+    info = {}
+    for header in info_headers:
+        info[header] = metadata[header]
+
+    return table, info
+
+def create_manual_marking(metadata, out_path=None):
+    """
+    Return:
+        string, full contents of manual marking
+    """
+
+    s = f'''
+Total Fish       =      {metadata["TOTAL_FISH"]}
+Upstream         =      {metadata["UPSTREAM_FISH"]}
+Downstream       =      {metadata["DOWNSTREAM_FISH"]}
+??               =      {metadata["NONDIRECTIONAL_FISH"]}
+
+Total Frames     = {metadata["TOTAL_FRAMES"]}
+Expected Frames  = {metadata["EXPECTED_FRAMES"]}
+Total Time       = {metadata["TOTAL_TIME"]}
+Expected Time    = {metadata["EXPECTED_TIME"]}
+
+Upstream Motion  = {metadata["UPSTREAM_MOTION"]}
+
+Count  File Name:   {metadata["COUNT_FILE_NAME"]}
+Editor ID          = {metadata["EDITOR_ID"]}
+Intensity          = {metadata["INTENSITY"]}
+Threshold          = {metadata["THRESHOLD"]}
+Window Start       = {metadata["WINDOW_START"]:.2f}
+Window End         = {metadata["WINDOW_END"]:.2f}
+Water Temperature  = {metadata["WATER_TEMP"]}
+
+
+*** Manual Marking (Manual Sizing:  Q = Quality, N = Repeat Count) ***
+
+File  Total  Frame#  Dir   R (m)   Theta   L(cm)  dR(cm)    L/dR  Aspect    Time      Date      Latitude            Longitude               Pan      Tilt     Roll   Species                           Motion            Q   N  Comment
+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+'''
+    
+    for entry in metadata["FISH"]:
+        s += f'{entry["FILE"]:>4}  {entry["TOTAL"]:>5}  {entry["FRAME_NUM"]:>6}  {entry["DIR"]:>3}  {entry["R"]:>6.2f}  {entry["THETA"]:>6.1f}  {entry["L"]:>6.1f}  {entry["DR"]:>6.1f}  {entry["LDR"]:>6.2f}  {entry["ASPECT"]:>6.1f}  {entry["TIME"]:>8}  {entry["DATE"]:>10}  {entry["LATITUDE"]:>17}  {entry["LONGITUDE"]:>18}  {entry["PAN"]:>7.2f}  {entry["TILT"]:>7.2f}  {entry["ROLL"]:>7.2f}  {entry["SPECIES"]:>8}  {entry["MOTION"]:>37}  {entry["Q"]:>5}  {entry["N"]:>2}  {entry["COMMENT"]}\n'
+
+    s += f'''
+*** Source File Key ***
+
+1.	Source File Name:   {metadata["FILE_NAME"]}
+    Source File Date:   {metadata["DATE"]}
+    Source File Start:  {metadata["START"]}
+    Source File End:    {metadata["END"]}
+
+Settings
+Upstream: {metadata["UPSTREAM_MOTION"]}
+Default Mark Direction: Upstream
+Editor ID: {metadata["EDITOR_ID"]}
+Show Marks: ??
+Show marks for ?? seconds
+Loop for ?? seconds
+'''
+    if out_path:
+        with open(out_path, 'w') as f:
+            f.write(s)
+    return s

aws_handler.py

@@ -0,0 +1,27 @@
+import logging
+import boto3
+from botocore.exceptions import ClientError
+import os
+
+
+def upload_file(file_name, bucket, object_name=None):
+    """Upload a file to an S3 bucket
+
+    :param file_name: File to upload
+    :param bucket: Bucket to upload to
+    :param object_name: S3 object name. If not specified then file_name is used
+    :return: True if file was uploaded, else False
+    """
+
+    # If S3 object_name was not specified, use file_name
+    if object_name is None:
+        object_name = os.path.basename(file_name)
+
+    # Upload the file
+    s3_client = boto3.client('s3')
+    try:
+        response = s3_client.upload_file(file_name, bucket, object_name)
+    except ClientError as e:
+        logging.error(e)
+        return False
+    return True

dataloader.py

@@ -0,0 +1,367 @@
+import project_path
+
+import os
+import cv2
+import numpy as np
+import json
+from threading import Lock
+import struct
+from contextlib import contextmanager
+import torch
+from torch.utils.data import Dataset
+
+# assumes yolov5 on sys.path
+from lib.yolov5.utils.general import xyxy2xywh
+from lib.yolov5.utils.augmentations import letterbox
+from lib.yolov5.utils.dataloaders import create_dataloader as create_yolo_dataloader
+
+from pyDIDSON import pyDIDSON
+from aris import ImageData
+
+# use this flag to test the difference between direct ARIS dataloading and
+# using the jpeg compressed version. very slow. not much difference observed.
+TEST_JPG_COMPRESSION = False
+
+
+# # # # # #
+# Factory(ish) methods for DataLoader creation. Easy entry points to this module.
+# # # # # #
+
+def create_dataloader_aris(aris_filepath, beam_width_dir, annotations_file, batch_size=32, stride=64, pad=0.5, img_size=896, rank=-1, world_size=1, workers=0, 
+                      disable_output=False, cache_bg_frames=False):
+    """
+    Get a PyTorch Dataset and DataLoader for ARIS files with (optional) associated fisheye-formatted labels.
+    """
+    # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
+    # this is a no-op for a single-gpu machine
+    with torch_distributed_zero_first(rank):
+        dataset = YOLOARISBatchedDataset(aris_filepath, beam_width_dir, annotations_file, stride, pad, img_size, 
+                                         disable_output=disable_output, cache_bg_frames=cache_bg_frames)
+
+    
+    batch_size = min(batch_size, len(dataset))
+    nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers])  # number of workers
+    
+    if not disable_output:
+        print("dataset size", len(dataset))
+        print("dataset shape", dataset.shape)
+        print("Num workers", nw)
+#     sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None # if extending to multi-GPU inference, will need this
+    dataloader = torch.utils.data.dataloader.DataLoader(dataset, 
+                                                        batch_size=None,
+                                                        sampler=OnePerBatchSampler(data_source=dataset, batch_size=batch_size),
+                                                        num_workers=nw,
+                                                        pin_memory=True,
+                                                        collate_fn=collate_fn)
+    return dataloader, dataset
+
+def create_dataloader_frames(frames_path, batch_size=32, model_stride_max=32, 
+                             pad=0.5, img_size=896, rank=-1, world_size=1, workers=0, disable_output=False):
+    """
+    Create a DataLoader for a directory of frames without labels.
+    
+    Args:
+        model_stride_max: use model.stride.max()
+    """
+    
+    gs = max(int(model_stride_max), 32)  # grid size (max stride)
+    return create_yolo_dataloader(frames_path, img_size, batch_size, gs, single_cls=False, augment=False,
+                                       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):
+        """
+        A PyTorch Dataset class for loading an ARIS file and (optional) associated fisheye-format labels. 
+        This class handles the ARIS frame extraction and 3-channel representation generation.
+        
+        It is called a "BatchedDataset" because it loads contiguous frames in self.batch_size chunks. 
+        ** The PyTorch sampler must be aware of this!! ** Use the OnePerBatchSampler in this module when using this Dataset.
+        
+        Args:
+            cache_bg_frames: keep the frames used for bg subtraction stored in memory. careful of memory issues. only recommended
+                            for small values of num_frames_bg_subtract
+        """
+        # open ARIS data stream - TODO: make sure this is one per worker
+        self.data = open(aris_filepath, 'rb')
+        self.data_lock = Lock()
+        self.beam_width_dir = beam_width_dir
+        self.disable_output = disable_output
+        self.aris_filepath = aris_filepath
+        self.cache_bg_frames = cache_bg_frames
+        
+        # get header info
+        self.didson = pyDIDSON(self.aris_filepath, beam_width_dir=beam_width_dir)
+        self.xdim = self.didson.info['xdim']
+        self.ydim = self.didson.info['ydim']
+
+        # disable automatic batching - do it ourselves, reading batch_size frames from
+        # the ARIS file at a time
+        self.batch_size = batch_size
+        
+        # load fisheye annotations
+        if annotations_file is None:
+            if not self.disable_output:
+                print("Loading file with no labels.")
+            self.start_frame = self.didson.info['startframe']
+            self.end_frame = self.didson.info['endframe'] or self.didson.info['numframes']
+            self.labels = None
+        else:
+            self._load_labels(annotations_file)
+            
+        # intiialize the background subtraction
+        self.num_frames_bg_subtract = num_frames_bg_subtract
+        self._init_bg_frame()
+        
+    def _init_bg_frame(self):
+        """
+        Intialize bg frame for bg subtraction.
+        Uses min(self.num_frames_bg_subtract, total_frames) frames to do mean subtraction.
+        Caches these frames in self.extracted_frames for reuse.
+        """
+        # ensure the number of frames used is a multiple of self.batch_size so we can cache them and retrieve full batches
+        # add 1 extra frame to be used for optical flow calculation
+        num_frames_bg = min(self.end_frame - self.start_frame, self.num_frames_bg_subtract // self.batch_size * self.batch_size + 1)
+        
+        if not self.disable_output:
+            print("Initializing mean frame for background subtraction using", num_frames_bg, "frames...")
+        frames_for_bg_subtract = self.didson.load_frames(start_frame=self.start_frame, end_frame=self.start_frame + num_frames_bg)
+
+        ### NEW WAY ###
+        # save memory (and time?) by computing these in a streaming fashion vs. in a big batch
+        self.mean_blurred_frame = np.zeros([self.ydim, self.xdim], dtype=np.float32)
+        max_blurred_frame = np.zeros([self.ydim, self.xdim], dtype=np.float32)
+        for i in range(frames_for_bg_subtract.shape[0]):
+            blurred = cv2.GaussianBlur(
+                frames_for_bg_subtract[i],
+                (5,5),
+                0)
+            self.mean_blurred_frame += blurred
+            max_blurred_frame = np.maximum(max_blurred_frame, np.abs(blurred))
+        self.mean_blurred_frame /= frames_for_bg_subtract.shape[0]
+        max_blurred_frame -= self.mean_blurred_frame
+        self.mean_normalization_value = np.max(max_blurred_frame)
+        
+        # cache these for later
+        self.extracted_frames = []
+        
+        # Because of the optical flow computation, we only go to end_frame - 1
+        next_blur = None
+        for i in range(len(frames_for_bg_subtract) - 1):
+            if next_blur is None:
+                this_blur = ((cv2.GaussianBlur(frames_for_bg_subtract[i], (5,5), 0) - self.mean_blurred_frame) / self.mean_normalization_value + 1) / 2
+            else:
+                this_blur = next_blur
+            next_blur = ((cv2.GaussianBlur(frames_for_bg_subtract[i+1], (5,5), 0) - self.mean_blurred_frame) / self.mean_normalization_value + 1) / 2
+            frame_image = np.dstack([frames_for_bg_subtract[i], 
+                                     this_blur * 255, 
+                                     np.abs(next_blur - this_blur) * 255]).astype(np.uint8, copy=False)
+
+            if TEST_JPG_COMPRESSION:
+                from PIL import Image
+                import os
+                Image.fromarray(frame_image).save(f"tmp{i}.jpg", quality=95)
+                frame_image = cv2.imread(f"tmp{i}.jpg")[:, :, ::-1] # BGR to RGB
+                os.remove(f"tmp{i}.jpg")
+                
+            if self.cache_bg_frames:
+                self.extracted_frames.append(frame_image)
+                
+        if not self.disable_output:
+            print("Done initializing background frame.")
+        
+    def _load_labels(self, fisheye_json):
+        """Load labels from a fisheye-formatted json file into self.labels in normalized
+        xywh format.
+        """
+        js = json.load(open(fisheye_json, 'r'))
+        labels = []
+
+        for frame in js['frames']:
+
+            l = []
+            for fish in frame['fish']:
+                x, y, w, h = xyxy2xywh(fish['bbox'])
+                cx = x + w/2.0
+                cy = y + h/2.0
+                # Each row is `class x_center y_center width height` format. (Normalized)
+                l.append([0, cx, cy, w, h])
+
+            l = np.array(l, dtype=np.float32)
+            if len(l) == 0:
+                l = np.zeros((0, 5), dtype=np.float32)
+
+            labels.append(l)
+
+        self.labels = labels
+        self.start_frame = js['start_frame']
+        self.end_frame = js['end_frame']
+
+    def __len__(self):
+        # account for optical flow - we can't do the last frame
+        return self.end_frame - self.start_frame - 1
+
+    def _postprocess(self, frame_images, frame_labels):
+        raise NotImplementedError
+    
+    def __getitem__(self, idx):
+        """
+        Return a numpy array representing this batch of frames and labels according to pyARIS frame extraction logic.
+        This class returns a full batch rather than just 1 example, assuming a OnePerBatchSampler is used.
+        """
+        final_idx = min(idx+self.batch_size, len(self))
+        frame_labels = self.labels[idx:final_idx] if self.labels else None
+        
+        # see if we have already cached this from bg subtraction
+        # assumes len(self.extracted_frames) is a multiple of self.batch_size
+        if idx+1 < len(self.extracted_frames):
+            return self._postprocess(self.extracted_frames[idx:final_idx], frame_labels)
+        else:
+            frames = self.didson.load_frames(start_frame=self.start_frame+idx, end_frame=self.start_frame + final_idx + 1)
+            blurred_frames = frames.astype(np.float32)
+            for i in range(frames.shape[0]):
+                blurred_frames[i] = cv2.GaussianBlur(
+                    blurred_frames[i],
+                    (5,5),
+                    0
+                )
+            blurred_frames -= self.mean_blurred_frame
+            blurred_frames /= self.mean_normalization_value
+            blurred_frames += 1
+            blurred_frames /= 2
+            
+            frame_images = np.stack([ frames[:-1], blurred_frames[:-1] * 255, np.abs(blurred_frames[1:] - blurred_frames[:-1]) * 255 ], axis=-1).astype(np.uint8, copy=False)
+            
+            if TEST_JPG_COMPRESSION:
+                from PIL import Image
+                import os
+                new_frame_images = []
+                for image in frame_images:
+                    Image.fromarray(image).save(f"tmp{idx}.jpg", quality=95)
+                    image = cv2.imread(f"tmp{idx}.jpg")[:, :, ::-1] # BGR to RGB
+                    os.remove(f"tmp{idx}.jpg")
+                    new_frame_images.append(image)
+                frame_images = new_frame_images
+            
+            return self._postprocess(frame_images, frame_labels)
+        
+class YOLOARISBatchedDataset(ARISBatchedDataset):
+    """An ARIS Dataset that works with YOLOv5 inference."""
+    
+    def __init__(self, aris_filepath, beam_width_dir, annotations_file, stride=64, pad=0.5, img_size=896, batch_size=32, 
+                 disable_output=False, cache_bg_frames=False):
+        super().__init__(aris_filepath, beam_width_dir, annotations_file, batch_size, disable_output=disable_output, cache_bg_frames=cache_bg_frames)
+        
+        # compute shapes for letterbox
+        aspect_ratio = self.ydim / self.xdim
+        if aspect_ratio < 1:
+            shape = [aspect_ratio, 1]
+        elif aspect_ratio > 1:
+            shape = [1, 1 / aspect_ratio]
+        self.original_shape = (self.ydim, self.xdim)
+        self.shape = np.ceil(np.array(shape) * img_size / stride + pad).astype(int) * stride
+
+    @classmethod
+    def load_image(cls, img, img_size=896):
+        """Loads and resizes 1 image from dataset, returns img, original hw, resized hw.
+        Modified from ScaledYOLOv4.datasets.load_image()
+        """
+        h0, w0 = img.shape[:2]  # orig hw
+        r = img_size / max(h0, w0)  # resize image to img_size
+        if r != 1:  # always resize down, only resize up if training with augmentation
+            interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
+            img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
+        return img, (h0, w0), img.shape[:2]  # img, hw_original, hw_resized
+        
+    def _postprocess(self, frame_images, frame_labels):
+        """
+        Return a batch of data in the format used by ScaledYOLOv4.
+        That is, a list of tuples, on tuple per image in the batch:
+            [
+                (img ->torch.Tensor,
+                labels ->torch.Tensor,
+                shapes ->tuple describing image original dimensions and scaled/padded dimensions
+                ),
+                ...
+            ]
+        """
+        outputs = []
+        frame_labels = frame_labels or [ None for _ in frame_images ]
+        for image, x in zip(frame_images, frame_labels):
+            img, (h0, w0), (h, w) = self.load_image(image)
+
+            # Letterbox
+            img, ratio, pad = letterbox(img, self.shape, auto=False, scaleup=False)
+            shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling
+
+            img = img.transpose(2, 0, 1) # to -> C x H x W
+            img = np.ascontiguousarray(img)
+
+            # Load labels
+            # Convert from normalized xywh to pixel xyxy format in order to add padding from letterbox
+            labels = []
+            if x is not None and x.size > 0:
+                labels = x.copy()
+                labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0]  # pad width
+                labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1]  # pad height
+                labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0]
+                labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1]
+
+            # convert back to normalized xywh with padding
+            nL = len(labels)  # number of labels
+            labels_out = torch.zeros((nL, 6))
+            if nL:
+                labels[:, 1:5] = xyxy2xywh(labels[:, 1:5])  # convert xyxy to xywh
+                labels[:, [2, 4]] /= img.shape[1]  # normalized height 0-1
+                labels[:, [1, 3]] /= img.shape[2]  # normalized width 0-1
+                labels_out[:, 1:] = torch.from_numpy(labels)
+            
+            outputs.append( (torch.from_numpy(img), labels_out, shapes) )
+            
+        return outputs
+    
+@contextmanager
+def torch_distributed_zero_first(local_rank: int):
+    """
+    Decorator to make all processes in distributed training wait for each local_master to do something.
+    """
+    if local_rank not in [-1, 0]:
+        torch.distributed.barrier()
+    yield
+    if local_rank == 0:
+        torch.distributed.barrier()
+        
+class OnePerBatchSampler(torch.utils.data.Sampler):
+    """Yields the first index of each batch, given a batch size.
+    In other words, returns multiples of self.batch_size up to the size of the Dataset.
+    This is a workaround for Pytorch's standard batch creation that allows us to manually
+    select contiguous segments of an ARIS clip for each batch.
+    """
+
+    def __init__(self, data_source, batch_size):
+        self.data_source = data_source
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        idxs = [i*self.batch_size for i in range(len(self))]
+        return iter(idxs)
+
+    def __len__(self):
+        return len(self.data_source) // self.batch_size
+    
+def collate_fn(batch):
+    """See YOLOv5.utils.datasets.collate_fn"""
+    if not len(batch):
+        print("help!")
+        print(batch)
+        
+    img, label, shapes = zip(*batch) # transposed
+    for i, l in enumerate(label):
+            l[:, 0] = i  # add target image index for build_targets()
+    return torch.stack(img, 0), torch.cat(label, 0), shapes

file_reader.py

@@ -0,0 +1,281 @@
+"""gr.File() component"""
+
+from __future__ import annotations
+
+import tempfile
+import warnings
+from pathlib import Path
+from typing import Any, Callable, Literal
+
+from gradio_client import utils as client_utils
+from gradio_client.documentation import document, set_documentation_group
+from gradio_client.serializing import FileSerializable
+
+from gradio import utils
+from gradio.components.base import IOComponent, _Keywords
+from gradio.deprecation import warn_deprecation
+from gradio.events import (
+    Changeable,
+    Clearable,
+    EventListenerMethod,
+    Selectable,
+    Uploadable,
+)
+
+set_documentation_group("component")
+
+
+@document()
+class File(
+    Changeable,
+    Selectable,
+    Clearable,
+    Uploadable,
+    IOComponent,
+    FileSerializable,
+):
+    """
+    Creates a file component that allows uploading generic file (when used as an input) and or displaying generic files (output).
+    Preprocessing: passes the uploaded file as a {tempfile._TemporaryFileWrapper} or {List[tempfile._TemporaryFileWrapper]} depending on `file_count` (or a {bytes}/{List{bytes}} depending on `type`)
+    Postprocessing: expects function to return a {str} path to a file, or {List[str]} consisting of paths to files.
+    Examples-format: a {str} path to a local file that populates the component.
+    Demos: zip_to_json, zip_files
+    """
+
+    def __init__(
+        self,
+        value: str | list[str] | Callable | None = None,
+        *,
+        file_count: Literal["single", "multiple", "directory"] = "single",
+        file_types: list[str] | None = None,
+        type: Literal["file", "binary"] = "file",
+        label: str | None = None,
+        every: float | None = None,
+        show_label: bool | None = None,
+        container: bool = True,
+        scale: int | None = None,
+        min_width: int = 160,
+        interactive: bool | None = None,
+        visible: bool = True,
+        elem_id: str | None = None,
+        elem_classes: list[str] | str | None = None,
+        **kwargs,
+    ):
+        """
+        Parameters:
+            value: Default file to display, given as str file path. If callable, the function will be called whenever the app loads to set the initial value of the component.
+            file_count: if single, allows user to upload one file. If "multiple", user uploads multiple files. If "directory", user uploads all files in selected directory. Return type will be list for each file in case of "multiple" or "directory".
+            file_types: List of file extensions or types of files to be uploaded (e.g. ['image', '.json', '.mp4']). "file" allows any file to be uploaded, "image" allows only image files to be uploaded, "audio" allows only audio files to be uploaded, "video" allows only video files to be uploaded, "text" allows only text files to be uploaded.
+            type: Type of value to be returned by component. "file" returns a temporary file object with the same base name as the uploaded file, whose full path can be retrieved by file_obj.name, "binary" returns an bytes object.
+            label: component name in interface.
+            every: If `value` is a callable, run the function 'every' number of seconds while the client connection is open. Has no effect otherwise. Queue must be enabled. The event can be accessed (e.g. to cancel it) via this component's .load_event attribute.
+            show_label: if True, will display label.
+            container: If True, will place the component in a container - providing some extra padding around the border.
+            scale: relative width compared to adjacent Components in a Row. For example, if Component A has scale=2, and Component B has scale=1, A will be twice as wide as B. Should be an integer.
+            min_width: minimum pixel width, will wrap if not sufficient screen space to satisfy this value. If a certain scale value results in this Component being narrower than min_width, the min_width parameter will be respected first.
+            interactive: if True, will allow users to upload a file; if False, can only be used to display files. If not provided, this is inferred based on whether the component is used as an input or output.
+            visible: If False, component will be hidden.
+            elem_id: An optional string that is assigned as the id of this component in the HTML DOM. Can be used for targeting CSS styles.
+            elem_classes: An optional list of strings that are assigned as the classes of this component in the HTML DOM. Can be used for targeting CSS styles.
+        """
+        self.file_count = file_count
+        self.file_types = file_types
+        if file_types is not None and not isinstance(file_types, list):
+            raise ValueError(
+                f"Parameter file_types must be a list. Received {file_types.__class__.__name__}"
+            )
+        valid_types = [
+            "file",
+            "binary",
+            "bytes",
+        ]  # "bytes" is included for backwards compatibility
+        if type not in valid_types:
+            raise ValueError(
+                f"Invalid value for parameter `type`: {type}. Please choose from one of: {valid_types}"
+            )
+        if type == "bytes":
+            warn_deprecation(
+                "The `bytes` type is deprecated and may not work as expected. Please use `binary` instead."
+            )
+        if file_count == "directory" and file_types is not None:
+            warnings.warn(
+                "The `file_types` parameter is ignored when `file_count` is 'directory'."
+            )
+        self.type = type
+        self.select: EventListenerMethod
+        """
+        Event listener for when the user selects file from list.
+        Uses event data gradio.SelectData to carry `value` referring to name of selected file, and `index` to refer to index.
+        See EventData documentation on how to use this event data.
+        """
+        IOComponent.__init__(
+            self,
+            label=label,
+            every=every,
+            show_label=show_label,
+            container=container,
+            scale=scale,
+            min_width=min_width,
+            interactive=interactive,
+            visible=visible,
+            elem_id=elem_id,
+            elem_classes=elem_classes,
+            value=value,
+            **kwargs,
+        )
+
+    def get_config(self):
+        return {
+            "file_count": self.file_count,
+            "file_types": self.file_types,
+            "value": self.value,
+            "selectable": self.selectable,
+            **IOComponent.get_config(self),
+        }
+
+    @staticmethod
+    def update(
+        value: Any | Literal[_Keywords.NO_VALUE] | None = _Keywords.NO_VALUE,
+        label: str | None = None,
+        show_label: bool | None = None,
+        container: bool | None = None,
+        scale: int | None = None,
+        min_width: int | None = None,
+        interactive: bool | None = None,
+        visible: bool | None = None,
+    ):
+        return {
+            "label": label,
+            "show_label": show_label,
+            "container": container,
+            "scale": scale,
+            "min_width": min_width,
+            "interactive": interactive,
+            "visible": visible,
+            "value": value,
+            "__type__": "update",
+        }
+
+    def preprocess(
+        self, x: list[dict[str, Any]] | None
+    ) -> (
+        bytes
+        | tempfile._TemporaryFileWrapper
+        | list[bytes | tempfile._TemporaryFileWrapper]
+        | None
+    ):
+        """
+        Parameters:
+            x: List of JSON objects with filename as 'name' property and base64 data as 'data' property
+        Returns:
+            File objects in requested format
+        """
+        if x is None:
+            return None
+
+        def process_single_file(f) -> bytes | tempfile._TemporaryFileWrapper:
+            file_name, data, is_file = (
+                f["name"],
+                f["data"],
+                f.get("is_file", False),
+            )
+            if self.type == "file":
+                if is_file:
+                    path = self.make_temp_copy_if_needed(file_name)
+                else:
+                    data, _ = client_utils.decode_base64_to_binary(data)
+                    path = self.file_bytes_to_file(
+                        data, dir=self.DEFAULT_TEMP_DIR, file_name=file_name
+                    )
+                    path = str(utils.abspath(path))
+                    self.temp_files.add(path)
+
+                # Creation of tempfiles here
+                file = tempfile.NamedTemporaryFile(
+                    delete=False, dir=self.DEFAULT_TEMP_DIR
+                )
+                file.name = path
+                file.orig_name = file_name  # type: ignore
+                return file
+            elif (
+                self.type == "binary" or self.type == "bytes"
+            ):  # "bytes" is included for backwards compatibility
+                if is_file:
+                    with open(file_name, "rb") as file_data:
+                        return (file_name, file_data.read())
+                return (file_name, client_utils.decode_base64_to_binary(data)[0])
+            else:
+                raise ValueError(
+                    "Unknown type: "
+                    + str(self.type)
+                    + ". Please choose from: 'file', 'bytes'."
+                )
+
+        if self.file_count == "single":
+            if isinstance(x, list):
+                return process_single_file(x[0])
+            else:
+                return process_single_file(x)
+        else:
+            if isinstance(x, list):
+                return [process_single_file(f) for f in x]
+            else:
+                return process_single_file(x)
+
+    def postprocess(
+        self, y: str | list[str] | None
+    ) -> dict[str, Any] | list[dict[str, Any]] | None:
+        """
+        Parameters:
+            y: file path
+        Returns:
+            JSON object with key 'name' for filename, 'data' for base64 url, and 'size' for filesize in bytes
+        """
+        if y is None:
+            return None
+        if isinstance(y, list):
+            return [
+                {
+                    "orig_name": Path(file).name,
+                    "name": self.make_temp_copy_if_needed(file),
+                    "size": Path(file).stat().st_size,
+                    "data": None,
+                    "is_file": True,
+                }
+                for file in y
+            ]
+        else:
+            d = {
+                "orig_name": Path(y).name,
+                "name": self.make_temp_copy_if_needed(y),
+                "size": Path(y).stat().st_size,
+                "data": None,
+                "is_file": True,
+            }
+            return d
+
+    def as_example(self, input_data: str | list | None) -> str:
+        if input_data is None:
+            return ""
+        elif isinstance(input_data, list):
+            return ", ".join([Path(file).name for file in input_data])
+        else:
+            return Path(input_data).name
+
+    def api_info(self) -> dict[str, dict | bool]:
+        if self.file_count == "single":
+            return self._single_file_api_info()
+        else:
+            return self._multiple_file_api_info()
+
+    def serialized_info(self):
+        if self.file_count == "single":
+            return self._single_file_serialized_info()
+        else:
+            return self._multiple_file_serialized_info()
+
+    def example_inputs(self) -> dict[str, Any]:
+        if self.file_count == "single":
+            return self._single_file_example_inputs()
+        else:
+            return self._multiple_file_example_inputs()

inference.py

@@ -0,0 +1,167 @@
+import project_path
+
+import torch
+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
+###
+
+def norm(bbox, w, h):
+    """
+    Normalize a bounding box.
+    Args:
+        bbox: list of length 4. Can be [x,y,w,h] or [x0,y0,x1,y1]
+        w: image width
+        h: image height
+    """
+    bb = bbox.copy()
+    bb[0] /= w
+    bb[1] /= h
+    bb[2] /= w
+    bb[3] /= h
+    return bb
+
+def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None, weights=WEIGHTS):
+    
+    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)
+
+    return results
+    
+
+def setup_model(weights_fp=WEIGHTS, imgsz=896, batch_size=32):
+    if torch.cuda.is_available():
+        device = select_device('0', batch_size=batch_size)
+    else:
+        print("CUDA not available. Using CPU inference.")
+        device = select_device('cpu', batch_size=batch_size)
+    
+    # Setup model for inference
+    model = attempt_load(weights_fp, device=device)
+    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)
+    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
+    
+    return model, device
+                       
+def do_detection(dataloader, model, device, 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, "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
+            for si, pred in enumerate(output):
+                # 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
+                
+                # 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]:
+                    do_norm = partial(norm, w=shapes[si][0][1], h=shapes[si][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)
+                all_preds[frame_num] = boxes
+
+            pbar.update(1*batch_size)
+         
+    return all_preds
+
+def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None): 
+
+    if (gp): gp(0, "Tracking...")
+
+    # Initialize tracker
+    clip_info = {
+        'start_frame': 0,
+        'end_frame': len(all_preds),
+        '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:
+        for i, key in enumerate(sorted(all_preds.keys())):
+            if gp: gp(i / len(all_preds), pbar.__str__())
+            boxes = all_preds[key]
+            if boxes is not None:
+                tracker.update(boxes)
+            else:
+                tracker.update()
+            pbar.update(1)
+    json_data = tracker.finalize(min_length=min_length)
+
+    return json_data
+
+
+@patch('json.encoder.c_make_encoder', None)
+def json_dump_round_float(some_object, out_path, num_digits=4):
+    """Write a json file to disk with a specified level of precision.
+    See: https://gist.github.com/Sukonnik-Illia/ed9b2bec1821cad437d1b8adb17406a3
+    """
+    # saving original method
+    of = json.encoder._make_iterencode
+    def inner(*args, **kwargs):
+        args = list(args)
+        # fifth argument is float formater which will we replace
+        fmt_str = '{:.' + str(num_digits) + 'f}'
+        args[4] = lambda o: fmt_str.format(o)
+        return of(*args, **kwargs)
+    
+    with patch('json.encoder._make_iterencode', wraps=inner):
+        return json.dump(some_object, open(out_path, 'w'), indent=2)

main.py

@@ -0,0 +1,77 @@
+import os
+import torch
+from zipfile import ZipFile
+
+from aris import create_manual_marking, BEAM_WIDTH_DIR, create_metadata_dictionary, prep_for_mm
+from dataloader import create_dataloader_aris
+from inference import do_full_inference, json_dump_round_float
+from visualizer import generate_video_batches
+
+WEIGHTS = 'models/v5m_896_300best.pt'
+
+def predict_task(filepath, weights=WEIGHTS, gradio_progress=None):
+    """
+    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.
+    """
+    if (gradio_progress): gradio_progress(0, "In task...")
+    print("Cuda available in task?", torch.cuda.is_available())
+    
+    print(filepath)
+    dirname = os.path.dirname(filepath)
+    filename = os.path.basename(filepath).replace(".aris","").replace(".ddf","")
+    results_filepath = os.path.join(dirname, f"{filename}_results.json")
+    marking_filepath = os.path.join(dirname, f"{filename}_marking.txt")
+    video_filepath = os.path.join(dirname, f"{filename}_results.mp4")
+    zip_filepath = os.path.join(dirname, f"{filename}_results.zip")
+    os.makedirs(dirname, exist_ok=True)
+
+    # create dataloader
+    if (gradio_progress): gradio_progress(0, "Initializing Dataloader...")
+    dataloader, dataset = create_dataloader_aris(filepath, BEAM_WIDTH_DIR, None)
+
+    # extract aris/didson info. didson does not yet have pixel-meter info
+    if ".ddf" in filepath:
+        image_meter_width = -1
+        image_meter_height = -1
+    else:
+        image_meter_width = dataset.didson.info['xdim'] * dataset.didson.info['pixel_meter_width']
+        image_meter_height = dataset.didson.info['ydim'] * dataset.didson.info['pixel_meter_height']
+    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)
+    
+    # write output to disk
+    json_dump_round_float(results, results_filepath)
+    
+    metadata = None
+    if dataset.didson.info['version'][3] == 5: # ARIS only
+        metadata = create_metadata_dictionary(filepath, results_filepath)
+        create_manual_marking(metadata, out_path=marking_filepath)
+    
+    # generate a video with tracking results
+    generate_video_batches(dataset.didson, results_filepath, frame_rate, video_filepath, 
+                   image_meter_width=image_meter_width, image_meter_height=image_meter_height, gp=gradio_progress)
+    
+    # zip up the results
+    with ZipFile(zip_filepath, 'w') as z:
+        for file in [results_filepath, marking_filepath, video_filepath, os.path.join(dirname, 'bg_start.jpg')]:
+            if os.path.exists(file):
+                z.write(file, arcname=os.path.basename(file))
+    
+    # release GPU memory 
+    torch.cuda.empty_cache()
+    
+    return metadata, results_filepath, zip_filepath, video_filepath, marking_filepath

project_path.py

@@ -0,0 +1,11 @@
+"""
+A kind of hacky way to get all the lib paths in order.
+"""
+
+import sys
+import os
+
+current_dir = os.path.dirname(os.path.realpath(__file__))
+for d in [current_dir, os.path.join(current_dir, "lib/fish_eye/"), os.path.join(current_dir, "lib/"), os.path.join(current_dir, "lib/yolov5/")]:
+    if d not in sys.path:
+        sys.path.append(d)

pyDIDSON.py

@@ -0,0 +1,495 @@
+"""
+Utilities to read and produce to-scale images from DIDSON and ARIS sonar files.
+
+Portions of this code were adapted from SoundMetrics MATLAB code.
+"""
+__version__ = 'b1.0.2'
+
+import contextlib
+import itertools
+from matplotlib.cm import get_cmap
+import numpy as np
+import os
+import pandas as pd
+from PIL import Image
+from shutil import make_archive, rmtree
+import struct
+from types import SimpleNamespace
+
+import lib.fish_eye.pyARIS as pyARIS
+from pyDIDSON_format import *
+
+
+class pyDIDSON:
+    def __init__(self, file, beam_width_dir='beam_widths', ixsize=-1):
+        """ Load header info from DIDSON file and precompute some warps.
+        
+        Parameters
+        ----------
+        file : file-like object, string, or pathlib.Path
+            The DIDSON or ARIS file to read.
+        beam_width_dir : string or pathlib.Path, optional
+            Location of ARIS beam width CSV files. Only used for ARIS files.
+        ixsize : int, optional
+            x-dimension width of output warped images to produce. Width is approximate for ARIS files and definite for
+            DIDSON. If not specified, the default for ARIS is determined by pyARIS and the default for DIDSON is 300.
+
+        Returns
+        -------
+        info : dict
+            Dictionary of extracted headers and computed sonar values.
+
+        """
+
+        if hasattr(file, 'read'):
+            file_ctx = contextlib.nullcontext(file)
+        else:
+            file_ctx = open(file, 'rb')
+
+        with file_ctx as fid:
+            assert fid.read(3) == b'DDF'
+
+            version_id = fid.read(1)[0]
+            print(f'Version {version_id}')
+
+            fid.seek(0)
+
+            info = {
+                'pydidson_version': __version__,
+            }
+            self.info = info
+
+            file_attributes, frame_attributes = {
+                0: NotImplementedError,
+                1: NotImplementedError,
+                2: NotImplementedError,
+                3: [file_attributes_3, frame_attributes_3],
+                4: [file_attributes_4, frame_attributes_4],
+                5: [file_attributes_5, frame_attributes_5],
+            }[version_id]
+
+            fileheaderformat = '=' + ''.join(file_attributes.values())
+            fileheadersize = struct.calcsize(fileheaderformat)
+            info.update(dict(zip(file_attributes.keys(), struct.unpack(fileheaderformat, fid.read(fileheadersize)))))
+
+            frameheaderformat = '=' + ''.join(frame_attributes.values())
+            frameheadersize = struct.calcsize(frameheaderformat)
+            info.update(dict(zip(frame_attributes.keys(), struct.unpack(frameheaderformat, fid.read(frameheadersize)))))
+
+            info.update({
+                'fileheaderformat': fileheaderformat,
+                'fileheadersize': fileheadersize,
+                'frameheaderformat': frameheaderformat,
+                'frameheadersize': frameheadersize,
+            })
+
+            if version_id == 0:
+                raise NotImplementedError
+            elif version_id == 1:
+                raise NotImplementedError
+            elif version_id == 2:
+                raise NotImplementedError
+            elif version_id == 3:
+                # Convert windowlength code to meters
+                info['windowlength'] = {
+                    0b00: [0.83, 2.5, 5, 10, 20, 40],  # DIDSON-S, Extended Windows
+                    0b01: [1.125, 2.25, 4.5, 9, 18, 36],  # DIDSON-S, Classic Windows
+                    0b10: [2.5, 5, 10, 20, 40, 70],  # DIDSON-LR, Extended Window
+                    0b11: [2.25, 4.5, 9, 18, 36, 72],  # DIDSON-LR, Classic Windows
+                }[info['configflags'] & 0b11][info['windowlength'] + 2 * (1 - info['resolution'])]
+
+                # Windowstart 1 to 31 times 0.75 (Lo) or 0.375 (Hi) or 0.419 for extended
+                info['windowstart'] = {
+                    0b0: 0.419 * info['windowstart'] * (2 - info['resolution']),  # meters for extended DIDSON
+                    0b1:
+                    0.375 * info['windowstart'] * (2 - info['resolution']),  # meters for standard or long range DIDSON
+                }[info['configflags'] & 0b1]
+
+                info['halffov'] = 14.4
+            elif version_id == 4:
+                # Convert windowlength code to meters
+                info['windowlength'] = [1.25, 2.5, 5, 10, 20, 40][info['windowlength'] + 2 * (1 - info['resolution'])]
+
+                # Windowstart 1 to 31 times 0.75 (Lo) or 0.375 (Hi) or 0.419 for extended
+                info['windowstart'] = 0.419 * info['windowstart'] * (2 - info['resolution'])
+
+                info['halffov'] = 14.4
+            elif version_id == 5:  #ARIS
+                if info['pingmode'] in [1, 2]:
+                    BeamCount = 48
+                elif info['pingmode'] in [3, 4, 5]:
+                    BeamCount = 96
+                elif info['pingmode'] in [6, 7, 8]:
+                    BeamCount = 64
+                elif info['pingmode'] in [9, 10, 11, 12]:
+                    BeamCount = 128
+                else:
+                    raise
+
+                WinStart = info['samplestartdelay'] * 0.000001 * info['soundspeed'] / 2
+
+                info.update({
+                    'BeamCount': BeamCount,
+                    'WinStart': WinStart,
+                })
+
+                aris_frame = SimpleNamespace(**info)
+
+                beam_width_data, camera_type = pyARIS.load_beam_width_data(frame=aris_frame,
+                                                                           beam_width_dir=beam_width_dir)
+
+                # What is the meter resolution of the smallest sample?
+                min_pixel_size = pyARIS.get_minimum_pixel_meter_size(aris_frame, beam_width_data)
+
+                # What is the meter resolution of the sample length?
+                sample_length = aris_frame.sampleperiod * 0.000001 * aris_frame.soundspeed / 2
+
+                # Choose the size of a pixel (or hard code it to some specific value)
+                pixel_meter_size = max(min_pixel_size, sample_length)
+
+                # Determine the image dimensions
+                xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds(
+                    pixel_meter_size,
+                    aris_frame,
+                    beam_width_data,
+                    additional_pixel_padding_x=0,
+                    additional_pixel_padding_y=0)
+
+                if ixsize != -1:
+                    pixel_meter_size = pixel_meter_size * xdim / ixsize
+                    pixel_meter_size += 1e-5
+                    xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds(
+                        pixel_meter_size,
+                        aris_frame,
+                        beam_width_data,
+                        additional_pixel_padding_x=0,
+                        additional_pixel_padding_y=0)
+
+                read_rows, read_cols, write_rows, write_cols = pyARIS.compute_mapping_from_sample_to_image(
+                    pixel_meter_size, xdim, ydim, x_meter_start, y_meter_start, aris_frame, beam_width_data)
+
+                read_i = read_rows * info['numbeams'] + info['numbeams'] - read_cols - 1
+
+                pixel_meter_width = pixel_meter_size
+                pixel_meter_height = pixel_meter_size
+
+                info.update({
+                    'camera_type': camera_type,
+                    'min_pixel_size': min_pixel_size,
+                    'sample_length': sample_length,
+                    'x_meter_start': x_meter_start,
+                    'y_meter_start': y_meter_start,
+                    'x_meter_stop': x_meter_stop,
+                    'y_meter_stop': y_meter_stop,
+                    'beam_width_dir': os.path.abspath(beam_width_dir),
+                })
+            else:
+                raise
+
+            if version_id < 5:
+                info['xdim'] = 300 if ixsize == -1 else ixsize
+                ydim, xdim, write_rows, write_cols, read_i = self.__mapscan()
+
+                # widthscale meters/pixels
+                pixel_meter_width = 2 * (info['windowstart'] + info['windowlength']) * np.sin(np.radians(14.25)) / xdim
+                # heightscale meters/pixels
+                pixel_meter_height = ((info['windowstart'] + info['windowlength']) -
+                                      info['windowstart'] * np.cos(np.radians(14.25))) / ydim
+
+                pixel_meter_size = (pixel_meter_width + pixel_meter_height) / 2
+
+            self.write_rows = write_rows
+            self.write_cols = write_cols
+            self.read_i = read_i
+
+            info.update({
+                'xdim': xdim,
+                'ydim': ydim,
+                'pixel_meter_width': pixel_meter_width,
+                'pixel_meter_height': pixel_meter_height,
+                'pixel_meter_size': pixel_meter_size,
+            })
+
+            # Fix common but critical corruption errors
+            if info['startframe'] > 65535:
+                info['startframe'] = 0
+            if info['endframe'] > 65535:
+                info['endframe'] = 0
+
+            try:
+                info['filename'] = os.path.abspath(file_ctx.name)
+            except AttributeError:
+                info['filename'] = None
+
+            # Record the proportion of measurements that are present in the warp (increases as xdim increases)
+            info['proportion_warp'] = len(np.unique(read_i)) / (info['numbeams'] * info['samplesperchannel'])
+
+    def __lens_distortion(self, nbeams, theta):
+        """ Removes Lens distortion determined by empirical work at the barge.
+        
+        Parameters
+        ----------
+        nbeams : int
+            Number of sonar beams.
+        theta : (A,) ndarray
+            Angle of warp for each x index.
+
+        Returns
+        -------
+        beamnum : (A,) ndarray
+            Distortion-adjusted beam number for each theta.
+
+        """
+
+        factor, a = {
+            48: [1, [.0015, -0.0036, 1.3351, 24.0976]],
+            189: [4.026, [.0015, -0.0036, 1.3351, 24.0976]],
+            96: [1.012, [.0030, -0.0055, 2.6829, 48.04]],
+            381: [4.05, [.0030, -0.0055, 2.6829, 48.04]],
+        }[nbeams]
+
+        return np.rint(factor * (a[0] * theta**3 + a[1] * theta**2 + a[2] * theta + a[3]) + 1).astype(np.uint32)
+
+    def __mapscan(self):
+        """ Calculate warp mapping from raw to scale images.
+
+        Returns
+        -------
+        ydim : int
+            y-dimension of warped image.
+        xdim : int
+            x-dimension of warped image.
+        write_rows : (A,) ndarray, np.uint16
+            Row indices to write to warped image.
+        write_cols : (A,) ndarray, np.uint16
+            Column indices to write to warped image.
+        read_i : (A,) ndarray, np.uint32
+            Indices to read from raw sonar measurements.
+
+        """
+
+        xdim = self.info['xdim']
+        rmin = self.info['windowstart']
+        rmax = rmin + self.info['windowlength']
+        halffov = self.info['halffov']
+        nbeams = self.info['numbeams']
+        nbins = self.info['samplesperchannel']
+
+        degtorad = 3.14159 / 180  # conversion of degrees to radians
+        radtodeg = 180 / 3.14159  # conversion of radians to degrees
+
+        d2 = rmax * np.cos(
+            halffov * degtorad)  # see drawing (distance from point scan touches image boundary to origin)
+        d3 = rmin * np.cos(halffov * degtorad)  # see drawing (bottom of image frame to r,theta origin in meters)
+        c1 = (nbins - 1) / (rmax - rmin)  # precalcualtion of constants used in do loop below
+        c2 = (nbeams - 1) / (2 * halffov)
+
+        gamma = xdim / (2 * rmax * np.sin(halffov * degtorad))  # Ratio of pixel number to position in meters
+        ydim = int(np.fix(gamma * (rmax - d3) + 0.5))  # number of pixels in image in vertical direction
+        svector = np.zeros(xdim * ydim, dtype=np.uint32)  # make vector and fill in later
+        ix = np.arange(1, xdim + 1)  # pixels in x dimension
+        x = ((ix - 1) - xdim / 2) / gamma  # convert from pixels to meters
+
+        for iy in range(1, ydim + 1):
+            y = rmax - (iy - 1) / gamma  # convert from pixels to meters
+            r = np.sqrt(y**2 + x**2)  # convert to polar cooridinates
+            theta = radtodeg * np.arctan2(x, y)  # theta is in degrees
+            binnum = np.rint((r - rmin) * c1 + 1.5).astype(np.uint32)  # the rangebin number
+            beamnum = self.__lens_distortion(nbeams, theta)  # remove lens distortion using empirical formula
+
+            # find position in sample array expressed as a vector
+            # make pos = 0 if outside sector, else give it the offset in the sample array
+            pos = (beamnum > 0) * (beamnum <= nbeams) * (binnum > 0) * (binnum <= nbins) * (
+                (beamnum - 1) * nbins + binnum)
+            svector[(ix - 1) * ydim + iy - 1] = pos  # The offset in this array is the pixel offset in the image array
+            # The value at this offset is the offset in the sample array
+
+        svector = svector.reshape(xdim, ydim).T.flat
+        svectori = svector != 0
+
+        read_i = np.flipud(np.arange(nbins * nbeams, dtype=np.uint32).reshape(nbins,
+                                                                              nbeams).T).flat[svector[svectori] - 1]
+        write_rows, write_cols = np.unravel_index(np.where(svectori)[0], (ydim, xdim))
+        return ydim, xdim, write_rows.astype(np.uint16), write_cols.astype(np.uint16), read_i
+
+    def __FasterDIDSONRead(self, file, start_frame, end_frame):
+        """ Load raw frames from DIDSON.
+        
+        Parameters
+        ----------
+        file : file-like object, string, or pathlib.Path
+            The DIDSON or ARIS file to read.
+        info : dict
+            Dictionary of extracted headers and computed sonar values.
+        start_frame : int
+            Zero-indexed start of frame range (inclusive).
+        end_frame : int
+            End of frame range (exclusive).
+
+        Returns
+        -------
+        raw_frames : (end_frame - start_frame, framesize) ndarray, np.uint8
+            Extracted and flattened raw sonar measurements for frame range.
+
+        """
+
+        if hasattr(file, 'read'):
+            file_ctx = contextlib.nullcontext(file)
+        else:
+            file_ctx = open(file, 'rb')
+
+        with file_ctx as fid:
+            framesize = self.info['samplesperchannel'] * self.info['numbeams']
+            frameheadersize = self.info['frameheadersize']
+
+            fid.seek(self.info['fileheadersize'] + start_frame * (frameheadersize + framesize) + frameheadersize, 0)
+
+            return np.array([
+                np.frombuffer(fid.read(framesize + frameheadersize)[:framesize], dtype=np.uint8)
+                for _ in range(end_frame - start_frame)
+            ],
+                            dtype=np.uint8)
+
+    def load_frames(self, file=None, start_frame=-1, end_frame=-1):
+        """ Load and warp DIDSON frames into images.
+        
+        Parameters
+        ----------
+        file : file-like object, string, or pathlib.Path, optional
+            The DIDSON or ARIS file to read. Defaults to `filename` in `info`.
+        start_frame : int, optional
+            Zero-indexed start of frame range (inclusive). Defaults to the first available.
+        end_frame : int, optional
+            End of frame range (exclusive). Defaults to the last available frame.
+
+        Returns
+        -------
+        frames : (end_frame - start_frame, ydim, xdim) ndarray, np.uint8
+            Warped-to-scale sonar image tensor.
+
+        """
+        if file is None:
+            file = self.info['filename']
+
+        if hasattr(file, 'read'):
+            file_ctx = contextlib.nullcontext(file)
+        else:
+            file_ctx = open(file, 'rb')
+
+        with file_ctx as fid:
+            svector = None
+            if start_frame == -1:
+                start_frame = self.info['startframe']
+            if end_frame == -1:
+                end_frame = self.info['endframe'] or self.info['numframes']
+
+            data = self.__FasterDIDSONRead(fid, start_frame, end_frame)
+            frames = np.zeros((end_frame - start_frame, self.info['ydim'], self.info['xdim']), dtype=np.uint8)
+            frames[:, self.write_rows, self.write_cols] = data[:, self.read_i]
+            return frames
+
+    @staticmethod
+    def save_frames(path, frames, pad_zeros=False, multiprocessing=False, ydim=None, xdim=None, quality='web_high'):
+        """ Save frames as JPEG images.
+        
+        Parameters
+        ----------
+        path : string or pathlib.Path
+            Directory to output images to or zip file.
+        frames : (end_frame - start_frame, ydim, xdim) ndarray, np.uint8
+            Warped-to-scale sonar image tensor.
+        pad_zeros : bool, optional
+            If enabled adds appropriately padded zeros to filenames so alphabetic sort of images returns expected
+            ordering. Note that this option is turned off by default for compatibility with vatic.js which requires
+            that filenames are not padded.
+        multiprocessing : bool, optional
+            If enabled adds multi-process optimization for writing images.
+        ydim : int, optional
+            If provided resizes image to given ydim before saving.
+        xdim : int, optional
+            If provided resizes image to given xdim before saving.
+        quality : int or str
+            Either integer 1-100 or JPEG compression preset seen here:
+            https://github.com/python-pillow/Pillow/blob/master/src/PIL/JpegPresets.py
+
+        """
+
+        path = str(path)
+
+        to_zip = path.endswith('.zip')
+
+        if to_zip:
+            path = os.path.splitext(path)[0]
+
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+        if pad_zeros:
+            filename = f'{path}/{{:0{int(np.ceil(np.log10(len(frames))))}}}.jpg'
+        else:
+            filename = f'{path}/{{}}.jpg'
+
+        ydim = ydim or frames.shape[1]
+        xdim = xdim or frames.shape[2]
+
+        viridis = get_cmap()
+
+        def f(n):
+            Image.fromarray(viridis(n[1], bytes=True)[..., :3]).resize((xdim, ydim)).save(filename.format(n[0]),
+                                                                                          quality=quality)
+
+        ns = enumerate(frames)
+        if multiprocessing:
+            __mpmap(f, ns)
+        else:
+            list(map(f, ns))
+
+        if to_zip:
+            make_archive(path, 'zip', path)
+            rmtree(path)
+
+
+def __mpmap(func, iterable, processes=os.cpu_count() - 1, niceness=1, threading=False, flatten=False):
+    """ Helper function to add simple multiprocessing capabilities.
+    
+    Parameters
+    ----------
+    func : function
+        Function to be mapped.
+    iterable : iterable
+        Domain to be mapped over.
+    processes : int, optional
+        Number of processes to spawn. Default is one for all but one CPU core.
+    niceness : int, optional
+        Process niceness.
+    threading : bool, optional
+        If enabled replaces multiprocessing with multithreading
+    flatten : bool, optional
+        If enabled chains map output together before returning.
+        
+    Returns
+    -------
+    output : list
+        Image of mapped func over iterable.
+
+    """
+
+    import multiprocess as mp
+    import multiprocess.dummy
+
+    def initializer():
+        os.nice(niceness)
+
+    pool_class = mp.dummy.Pool if threading else mp.Pool
+
+    pool = pool_class(processes=processes, initializer=initializer)
+
+    out = pool.map(func, iterable)
+
+    if flatten:
+        out = list(itertools.chain.from_iterable(out))
+
+    pool.close()
+    pool.join()
+
+    return out

pyDIDSON_format.py

@@ -0,0 +1,364 @@
+"""
+DIDSON and ARIS file and frame header formats
+"""
+
+file_attributes_3 = {
+    'version': '4s',
+    'numframes': 'i',
+    'framerate': 'i',
+    'resolution': 'i',  # 0=lo 1 = Hi
+    'numbeams': 'i',  # 48 Lo 96 Hi for standard mode
+    'samplerate': 'f',
+    'samplesperchannel': 'i',
+    'receivergain': 'i',  # 0-40 dB
+    'windowstart': 'i',  # Windowstart 1 to 31 
+    'windowlength': 'i',  # Windowlength coded as 0 to 3 
+    'reverse': 'i',
+    'serialnumber': 'i',
+    'date': '32s',  # date file was made
+    'idstring': '256s',  # User supplied identification notes
+    'id1': 'i',  # four user supplied integers
+    'id2': 'i',
+    'id3': 'i',
+    'id4': 'i',
+    'startframe': 'i',  # used if this is a snippet file from source file
+    'endframe': 'i',  # Used if this is a snippet file from source file
+    'timelapse': 'i',  # Logic 0 or 1 (1 = timelapse active);
+    'recordInterval': 'i',  # Ask Bill
+    'radioseconds': 'i',  # Needed for timelapse -- ask Bill
+    'frameinterval': 'i',  # Interval between frames in time lapse
+    'userassigned': '136s',  # User assigned space
+}
+
+file_attributes_4 = {
+    'version': '4s',
+    'numframes': 'i',
+    'framerate': 'i',
+    'resolution': 'i',  # 0=lo 1 = Hi
+    'numbeams': 'i',  # 48 Lo 96 Hi for standard mode
+    'samplerate': 'f',
+    'samplesperchannel': 'i',
+    'receivergain': 'i',  # 0-40 dB
+    'windowstart': 'i',  # Windowstart 1 to 31 
+    'windowlength': 'i',  # Windowlength coded as 0 to 3 
+    'reverse': 'i',
+    'serialnumber': 'i',
+    'date': '32s',  # date file was made
+    'idstring': '256s',  # User supplied identification notes
+    'id1': 'i',  # four user supplied integers
+    'id2': 'i',
+    'id3': 'i',
+    'id4': 'i',
+    'startframe': 'i',  # used if this is a snippet file from source file
+    'endframe': 'i',  # Used if this is a snippet file from source file
+    'timelapse': 'i',  # Logic 0 or 1 (1 = timelapse active);
+    'recordInterval': 'i',  # Ask Bill
+    'radioseconds': 'i',  # Needed for timelapse -- ask Bill
+    'frameinterval': 'i',  # Interval between frames in time lapse
+    'userassigned': '136s',  # User assigned space
+}
+
+file_attributes_5 = {
+    'version': '4s',
+    'numframes': 'I',  # Total frames in file
+    'framerate': 'I',  # Initial recorded frame rate
+    'resolution': 'I',  # Non-zero if HF, zero if LF
+    'numbeams': 'I',  # ARIS 3000 = 128/64, ARIS 1800 = 96/48, ARIS 1200 = 48
+    'samplerate': 'f',  # 1/Sample Period
+    'samplesperchannel': 'I',  # Number of range samples in each beam
+    'receivergain': 'I',  # Relative gain in dB:  0 - 40
+    'windowstart': 'f',  # Image window start range in meters (code [0..31] in DIDSON)
+    'windowlength': 'f',  # Image window length in meters  (code [0..3] in DIDSON)
+    'reverse': 'I',  # Non-zero = lens down (DIDSON) or lens up (ARIS), zero = opposite
+    'serialnumber': 'I',  # Sonar serial number
+    'strdate': '32s',  # Date that file was recorded
+    'idstring': '256s',  # User input to identify file in 256 characters
+    'id1': 'i',  # User-defined integer quantity
+    'id2': 'i',  # User-defined integer quantity
+    'id3': 'i',  # User-defined integer quantity
+    'id4': 'i',  # User-defined integer quantity
+    'startframe': 'I',  # First frame number from source file (for DIDSON snippet files)
+    'endframe': 'I',  # Last frame number from source file (for DIDSON snippet files)
+    'timelapse': 'I',  # Non-zero indicates time lapse recording
+    'recordInterval': 'I',  # Number of frames/seconds between recorded frames
+    'radioseconds': 'I',  # Frames or seconds interval
+    'frameinterval': 'I',  # Record every Nth frame
+    'flags': 'I',  # See DDF_04 file format document
+    'auxflags': 'I',  # See DDF_04 file format document
+    'sspd': 'I',  # Sound velocity in water
+    'flags3d': 'I',  # See DDF_04 file format document
+    'softwareversion': 'I',  # DIDSON software version that recorded the file
+    'watertemperature': 'I',  # Water temperature code:  0 = 5-15C, 1 = 15-25C, 2 = 25-35C
+    'salinity': 'I',  # Salinity code:  0 = fresh, 1 = brackish, 2 = salt
+    'pulselength': 'I',  # Added for ARIS but not used
+    'txmode': 'I',  # Added for ARIS but not used
+    'versionfgpa': 'I',  # Reserved for future use
+    'versionpsuc': 'I',  # Reserved for future use
+    'thumbnailfi': 'I',  # Frame index of frame used for thumbnail image of file
+    'filesize': 'Q',  # Total file size in bytes
+    'optionalheadersize': 'Q',  # Reserved for future use
+    'optionaltailsize': 'Q',  # Reserved for future use
+    'versionminor': 'I',  # DIDSON_ADJUSTED_VERSION_MINOR
+    'largelens': 'I',  # Non-zero if telephoto lens (large lens, hi-res lens, big lens) is present
+    'userassigned': '568s',  # Free space for user
+}
+
+frame_attributes_3 = {
+    'framenumber': 'i',
+    'frametime': 'i',
+    'frametime2': 'i',
+    'version': '4s',
+    'status': 'i',
+    'year': 'i',
+    'month': 'i',
+    'day': 'i',
+    'hour': 'i',
+    'minute': 'i',
+    'second': 'i',
+    'hsecond': 'i',
+    'transmit': 'i',  # bit2 = 2.0 MHz, bit1 = Enable, bit0 = HF_MODE
+    'windowstart': 'i',  # This will be updated at the end of this routine
+    'windowlength': 'i',  # Add 2 if low resolution (index between 1 and 6)
+    'threshold': 'i',
+    'intensity': 'i',
+    'receivergain': 'i',
+    'degc1': 'i',
+    'degc2': 'i',
+    'humidity': 'i',
+    'focus': 'i',
+    'battery': 'i',
+    'status1': '16s',  # User defined and supplied
+    'status2': '8s',  # User defined and supplied
+    'panwcom': 'f',  # Return from Pan/Tilt if used when compass present
+    'tiltwcom': 'f',  # Return from Pan/Tilt if used when compass is present
+    'velocity': 'f',  # Platform variables supplied by user
+    'depth': 'f',
+    'altitude': 'f',
+    'pitch': 'f',
+    'pitchrate': 'f',
+    'roll': 'f',
+    'rollrate': 'f',
+    'heading': 'f',
+    'headingrate': 'f',
+    'sonarpan': 'f',
+    'sonartilt': 'f',  # Read from compass if used, Read from Pan/Tilt if used and no compass
+    'sonarroll': 'f',  # Read from compass if used, Read from Pan/Tilt if used and no compass
+    'latitude': 'd',
+    'longitude': 'd',
+    'sonarposition': 'f',
+    'configflags': 'i',  # bit0: 1=classic, 0=extended windows; bit1: 0=Standard, 1=LR
+    'userassigned': '60s',  # Free space for user
+}
+
+frame_attributes_4 = {
+    'framenumber': 'i',
+    'frametime': 'i',
+    'frametime2': 'i',
+    'version': '4s',
+    'status': 'i',
+    'year': 'i',
+    'month': 'i',
+    'day': 'i',
+    'hour': 'i',
+    'minute': 'i',
+    'second': 'i',
+    'hsecond': 'i',
+    'transmit': 'i',  # bit2 = 2.0 MHz, bit1 = Enable, bit0 = HF_MODE
+    'windowstart': 'i',  # This will be updated at the end of this routine
+    'windowlength': 'i',  # Add 2 if low resolution (index between 1 and 6)
+    'threshold': 'i',
+    'intensity': 'i',
+    'receivergain': 'i',
+    'degc1': 'i',
+    'degc2': 'i',
+    'humidity': 'i',
+    'focus': 'i',
+    'battery': 'i',
+    'status1': '16s',  # User defined and supplied
+    'status2': '8s',  # User defined and supplied
+    'panwcom': 'f',  # Return from Pan/Tilt if used when compass present
+    'tiltwcom': 'f',  # Return from Pan/Tilt if used when compass is present
+    'velocity': 'f',  # Platform variables supplied by user
+    'depth': 'f',
+    'altitude': 'f',
+    'pitch': 'f',
+    'pitchrate': 'f',
+    'roll': 'f',
+    'rollrate': 'f',
+    'heading': 'f',
+    'headingrate': 'f',
+    'sonarpan': 'f',
+    'sonartilt': 'f',  # Read from compass if used, Read from Pan/Tilt if used and no compass
+    'sonarroll': 'f',  # Read from compass if used, Read from Pan/Tilt if used and no compass
+    'latitude': 'd',
+    'longitude': 'd',
+    'sonarposition': 'f',
+    'configflags': 'i',  # bit0: 1=classic, 0=extended windows; bit1: 0=Standard, 1=LR
+    'userassigned': '828s',  # Move pointer to end of frame header of length 1024 bytes
+}
+
+frame_attributes_5 = {
+    'framenumber': 'I',
+    'frametime': 'Q',  # Recording timestamp
+    'version': '4s',
+    'status': 'I',
+    'sonartimestamp': 'Q',
+    'tsday': 'I',
+    'tshour': 'I',
+    'tsminute': 'I',
+    'tssecond': 'I',
+    'tshsecond': 'I',
+    'transmitmode': 'I',
+    'windowstart': 'f',
+    'windowlength': 'f',
+    'threshold': 'I',
+    'intensity': 'i',
+    'receivergain': 'I',
+    'degc1': 'I',
+    'degc2': 'I',
+    'humidity': 'I',
+    'focus': 'I',
+    'battery': 'I',
+    'uservalue1': 'f',
+    'uservalue2': 'f',
+    'uservalue3': 'f',
+    'uservalue4': 'f',
+    'uservalue5': 'f',
+    'uservalue6': 'f',
+    'uservalue7': 'f',
+    'uservalue8': 'f',
+    'velocity': 'f',
+    'depth': 'f',
+    'altitude': 'f',
+    'pitch': 'f',
+    'pitchrate': 'f',
+    'roll': 'f',
+    'rollrate': 'f',
+    'heading': 'f',
+    'headingrate': 'f',
+    'compassheading': 'f',
+    'compasspitch': 'f',
+    'compassroll': 'f',
+    'latitude': 'd',
+    'longitude': 'd',
+    'sonarposition': 'f',
+    'configflags': 'I',
+    'beamtilt': 'f',
+    'targetrange': 'f',
+    'targetbearing': 'f',
+    'targetpresent': 'I',
+    'firmwarerevision': 'I',
+    'flags': 'I',
+    'sourceframe': 'I',
+    'watertemp': 'f',
+    'timerperiod': 'I',
+    'sonarx': 'f',
+    'sonary': 'f',
+    'sonarz': 'f',
+    'sonarpan': 'f',
+    'sonartilt': 'f',
+    'sonarroll': 'f',
+    'panpnnl': 'f',
+    'tiltpnnl': 'f',
+    'rollpnnl': 'f',
+    'vehicletime': 'd',
+    'timeggk': 'f',
+    'dateggk': 'I',
+    'qualityggk': 'I',
+    'numsatsggk': 'I',
+    'dopggk': 'f',
+    'ehtggk': 'f',
+    'heavetss': 'f',
+    'yeargps': 'I',
+    'monthgps': 'I',
+    'daygps': 'I',
+    'hourgps': 'I',
+    'minutegps': 'I',
+    'secondgps': 'I',
+    'hsecondgps': 'I',
+    'sonarpanoffset': 'f',
+    'sonartiltoffset': 'f',
+    'sonarrolloffset': 'f',
+    'sonarxoffset': 'f',
+    'sonaryoffset': 'f',
+    'sonarzoffset': 'f',
+    'tmatrix': '64s',
+    'samplerate': 'f',
+    'accellx': 'f',
+    'accelly': 'f',
+    'accellz': 'f',
+    'pingmode': 'I',
+    'frequencyhilow': 'I',
+    'pulsewidth': 'I',
+    'cycleperiod': 'I',
+    'sampleperiod': 'I',
+    'transmitenable': 'I',
+    'framerate': 'f',
+    'soundspeed': 'f',
+    'samplesperbeam': 'I',
+    'enable150v': 'I',
+    'samplestartdelay': 'I',
+    'largelens': 'I',
+    'thesystemtype': 'I',
+    'sonarserialnumber': 'I',
+    'encryptedkey': 'Q',
+    'ariserrorflagsuint': 'I',
+    'missedpackets': 'I',
+    'arisappversion': 'I',
+    'available2': 'I',
+    'reorderedsamples': 'I',
+    'salinity': 'I',
+    'pressure': 'f',
+    'batteryvoltage': 'f',
+    'mainvoltage': 'f',
+    'switchvoltage': 'f',
+    'focusmotormoving': 'I',
+    'voltagechanging': 'I',
+    'focustimeoutfault': 'I',
+    'focusovercurrentfault': 'I',
+    'focusnotfoundfault': 'I',
+    'focusstalledfault': 'I',
+    'fpgatimeoutfault': 'I',
+    'fpgabusyfault': 'I',
+    'fpgastuckfault': 'I',
+    'cputempfault': 'I',
+    'psutempfault': 'I',
+    'watertempfault': 'I',
+    'humidityfault': 'I',
+    'pressurefault': 'I',
+    'voltagereadfault': 'I',
+    'voltagewritefault': 'I',
+    'focuscurrentposition': 'I',
+    'targetpan': 'f',
+    'targettilt': 'f',
+    'targetroll': 'f',
+    'panmotorerrorcode': 'I',
+    'tiltmotorerrorcode': 'I',
+    'rollmotorerrorcode': 'I',
+    'panabsposition': 'f',
+    'tiltabsposition': 'f',
+    'rollabsposition': 'f',
+    'panaccelx': 'f',
+    'panaccely': 'f',
+    'panaccelz': 'f',
+    'tiltaccelx': 'f',
+    'tiltaccely': 'f',
+    'tiltaccelz': 'f',
+    'rollaccelx': 'f',
+    'rollaccely': 'f',
+    'rollaccelz': 'f',
+    'appliedsettings': 'I',
+    'constrainedsettings': 'I',
+    'invalidsettings': 'I',
+    'enableinterpacketdelay': 'I',
+    'interpacketdelayperiod': 'I',
+    'uptime': 'I',
+    'arisappversionmajor': 'H',
+    'arisappversionminor': 'H',
+    'gotime': 'Q',
+    'panvelocity': 'f',
+    'tiltvelocity': 'f',
+    'rollvelocity': 'f',
+    'sentinel': 'I',
+    'userassigned': '292s',  # Free space for user
+}

requirements.txt

@@ -0,0 +1,79 @@
+flask
+requests
+black
+
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python>=4.1.2
+Pillow
+PyYAML>=5.3.1
+scipy>=1.4.1
+torch>=1.9.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+
+tensorboard>=2.4.1
+
+seaborn>=0.11.0
+pandas
+
+thop  # FLOPs computation
+
+pycocotools
+filterpy
+celery
+redis
+boto3
+
+
+# YOLOv5 requirements
+# Usage: pip install -r requirements.txt
+
+# Base ------------------------------------------------------------------------
+--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
+
+# 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

state_handler.py

@@ -0,0 +1,375 @@
+from aris import create_metadata_table
+
+example_metadata = {
+    "FILE_NAME":  "static/example_metadata/fisheye",
+    "FRAME_RATE":  6.548702716827393,
+    "UPSTREAM_FISH":  0,
+    "DOWNSTREAM_FISH":  0,
+    "NONDIRECTIONAL_FISH":  14,
+    "TOTAL_FISH":  14,
+    "TOTAL_FRAMES":  644,
+    "EXPECTED_FRAMES":  -1,
+    "TOTAL_TIME":  "0:01:38",
+    "EXPECTED_TIME":  "0:00:00",
+    "UPSTREAM_MOTION":  "Right To Left",
+    "COUNT_FILE_NAME":  "N/A",
+    "EDITOR_ID":  "N/A",
+    "INTENSITY":  "0.0 dB",
+    "THRESHOLD":  "0.0 dB",
+    "WINDOW_START":  1,
+    "WINDOW_END":  17,
+    "WATER_TEMP":  "13 degC",
+    "FISH":  [
+        {
+            "FILE":  1,
+            "TOTAL":  1,
+            "FRAME_NUM":  12,
+            "DIR":  " N/A",
+            "R":  13.403139282569885,
+            "THETA":  0.1706,
+            "L":  63.739999999999995,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:54:40",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  2,
+            "FRAME_NUM":  35,
+            "DIR":  " N/A",
+            "R":  13.206211097755432,
+            "THETA":  -9.1195,
+            "L":  73.33,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:54:44",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  3,
+            "FRAME_NUM":  122,
+            "DIR":  " N/A",
+            "R":  13.219339643409729,
+            "THETA":  -9.3961,
+            "L":  84.77,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:54:58",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  4,
+            "FRAME_NUM":  123,
+            "DIR":  "N/A",
+            "R":  12.996154367286682,
+            "THETA":  10.7991,
+            "L":  59.919999999999995,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:54:58",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  5,
+            "FRAME_NUM":  130,
+            "DIR":  " N/A",
+            "R":  12.484141086769105,
+            "THETA":  -8.2654,
+            "L":  70.89999999999999,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:54:59",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  6,
+            "FRAME_NUM":  218,
+            "DIR":  " N/A",
+            "R":  13.232468189064026,
+            "THETA":  -9.3961,
+            "L":  77.25999999999999,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:55:12",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  7,
+            "FRAME_NUM":  278,
+            "DIR":  " N/A",
+            "R":  13.967666745704651,
+            "THETA":  -12.8758,
+            "L":  37.51,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:55:22",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  8,
+            "FRAME_NUM":  302,
+            "DIR":  " N/A",
+            "R":  13.25872528037262,
+            "THETA":  -9.1195,
+            "L":  79.5,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:55:25",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  9,
+            "FRAME_NUM":  331,
+            "DIR":  " N/A",
+            "R":  13.25872528037262,
+            "THETA":  -9.1195,
+            "L":  80.67,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:55:30",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  10,
+            "FRAME_NUM":  450,
+            "DIR":  " N/A",
+            "R":  13.324368008644104,
+            "THETA":  -8.5535,
+            "L":  83.1,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:55:48",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  11,
+            "FRAME_NUM":  495,
+            "DIR":  " N/A",
+            "R":  13.481910556495666,
+            "THETA":  -9.1195,
+            "L":  86.39,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:55:55",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  12,
+            "FRAME_NUM":  526,
+            "DIR":  " N/A",
+            "R":  13.04866854990387,
+            "THETA":  10.5397,
+            "L":  55.37,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:56:00",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  13,
+            "FRAME_NUM":  538,
+            "DIR":  " N/A",
+            "R":  13.416267828224182,
+            "THETA":  -9.668,
+            "L":  82.38,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:56:02",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }, {
+            "FILE":  1,
+            "TOTAL":  14,
+            "FRAME_NUM":  624,
+            "DIR":  " N/A",
+            "R":  13.29811091733551,
+            "THETA":  -8.8385,
+            "L":  77.44,
+            "DR":  -1,
+            "LDR":  -1,
+            "ASPECT":  -1,
+            "TIME":  "11:56:16",
+            "DATE":  "2018-07-09",
+            "LATITUDE":  "N 00 d 0.00000 m",
+            "LONGITUDE":  "E 000 d 0.00000 m",
+            "PAN":  None,
+            "TILT":  None,
+            "ROLL":  0,
+            "SPECIES":  "Unknown",
+            "MOTION":  "Running <-->",
+            "Q":  -1,
+            "N":  -1,
+            "COMMENT":  ""
+        }
+    ],
+    "DATE":  "2018-07-09",
+    "START":  "11:54:39",
+    "END":  "11:56:18"
+}
+
+def load_example_result(result, table_headers, info_headers):
+    fish_table, fish_info = create_metadata_table(example_metadata, table_headers, info_headers)
+    result['path_zip'] = ["static/example/input_file_results.zip"]
+    result['path_video'] = ["static/example/input_file_results.mp4"]
+    result['path_json'] = ["static/example/input_file_results.json"]
+    result['path_marking'] = ["static/example/input_file_marking.txt"]
+    result['fish_table'] = [fish_table]
+    result['fish_info'] = [fish_info]
+
+
+def reset_state(result, state):
+
+    # Reset Result
+    result["path_video"] = []
+    result["path_zip"] = []
+    result["path_json"] = []
+    result["path_marking"] = []
+    result["fish_table"] = []
+    result["fish_info"] = []
+
+    # Reset State
+    state['files'] = []
+    state['index'] = 0
+    state['total'] = 0

uploader.py

@@ -0,0 +1,41 @@
+import project_path
+import os;
+from datetime import datetime;
+
+USER_DATA_DIR = "user_data/"
+
+
+def save_data(bytes, filename):
+    """Take a file and saved it to a new user_data folder"""
+    
+    dirname = create_data_dir()
+    filepath = os.path.join(dirname, filename) 
+    
+    assert bytes[0:3] == b'DDF'
+
+    with open(filepath, 'wb') as out:
+        out.write(bytes)
+        
+    # check this is actually a valid ARIS file to catch any malicious fish scientists
+    try:
+        with open(filepath, 'rb') as file:
+            assert file.read(3) == b'DDF'
+    except:
+        print("Bad file!", filepath)
+        return False, None, None
+    
+    return True, filepath, dirname
+    
+
+def allowed_file(filename):
+    """Only allow an ARIS file to be uploaded."""
+    return '.' in filename and \
+           filename.rsplit('.', 1)[1].lower() in ['aris', 'ddf']
+
+def create_data_dir():
+    """Create a (probably) unique directory for a task."""
+    dirname = os.path.join(USER_DATA_DIR, str(int(datetime.now().timestamp())))
+    if os.path.exists(dirname):
+        print("Warning,", dirname, "already exists.")
+    os.makedirs(dirname, exist_ok=True)
+    return dirname

visualizer.py

@@ -0,0 +1,191 @@
+import project_path
+
+import json
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+from lib.fish_eye.tracker import Tracker
+
+
+VERSION = "09/21"
+PRED_COLOR = (255, 0, 0) # blue
+WHITE = (255, 255, 255)
+BLACK = (0, 0, 0)
+BORDER_PAD = 3
+LINE_HEIGHT= 22
+VIDEO_HEIGHT = 700
+INFO_PANE_WIDTH = 100
+BOX_THICKNESS = 2
+FONT_SCALE = 0.65
+FONT_THICKNESS = 1
+
+def generate_video_batches(didson, preds_path, frame_rate, video_out_path, gp=None, image_meter_width=None, image_meter_height=None, batch_size=1000):
+    """Write a visualized video to video_out_path, given a didson object.
+    """
+    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
+    
+    with tqdm(total=end_frame, desc="Generating results video", ncols=0) as pbar:
+        for i in range(0, end_frame, batch_size):
+            batch_end = min(end_frame, i+batch_size)
+            frames = didson.load_frames(start_frame=i, end_frame=batch_end)
+            vid_frames, h, w = get_video_frames(frames, preds_path, frame_rate, image_meter_width, image_meter_height, start_frame=i)
+
+            if out is None:
+                out = cv2.VideoWriter(video_out_path, cv2.VideoWriter_fourcc(*'avc1'), frame_rate, [ int(1.5*w), h ] )
+
+            for j, frame in enumerate(vid_frames):
+                if gp: gp(( (i+j) / end_frame), 'Generating results video...')
+                out.write(frame)
+                pbar.update(1)
+
+            del frames
+            del vid_frames
+    
+    out.release()
+    
+def get_video_frames(frames, preds_path, frame_rate, image_meter_width=None, image_meter_height=None, start_frame=0):
+    """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'))
+    pred_lengths = { fish['id'] : "%.2fm" % fish['length'] for fish in preds['fish'] }
+    clip_pr_counts = Tracker.count_dirs(preds)
+    color_map = { fish['id'] : fish['color'] for fish in preds['fish'] }
+    
+    # filter JSON, if necessary (for shorter clips)
+    preds['frames'] = preds['frames'][start_frame:]
+    
+    vid_frames = []
+    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]):
+            frame_raw = cv2.resize(cv2.cvtColor(frame_raw, cv2.COLOR_GRAY2BGR), (w,h))
+            pred = preds['frames'][i]
+
+            for fish in pred['fish']:
+                xmin, ymin, xmax, ymax = fish['bbox']
+                left = int(round(xmin * w))
+                right = int(round(xmax * w))
+                top = int(round(ymin * h))
+                bottom = int(round(ymax * h))
+                fish_id = str(fish['fish_id'])
+                fish_len = pred_lengths[fish['fish_id']]
+                hexx = color_map[fish['fish_id']].lstrip('#')
+                color = tuple(int(hexx[i:i+2], 16) for i in (0, 2, 4))
+                draw_fish(frame_raw, left, right, top, bottom, color, fish_id, fish_len, anno_align="right")
+
+            # add axis to frame
+            frame_raw = add_axis(frame_raw, image_meter_width, image_meter_height)
+
+            # add info
+            frame_info_panel = np.zeros((h, int(0.5*w), 3)).astype(np.uint8)
+            frame = np.concatenate((frame_info_panel, frame_raw), axis=1)
+            cv2.putText(frame, f'VERSION: {VERSION}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*4), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+            cv2.putText(frame, f'Right count: {clip_pr_counts[0]}',  (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*3), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+            cv2.putText(frame, f'Left count: {clip_pr_counts[FONT_THICKNESS]}',  (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*2), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+            cv2.putText(frame, f'Other fish: {clip_pr_counts[2]}',  (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*1), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+    #         cv2.putText(frame, f'Upstream: {preds["upstream_direction"]}', (0, h-1-LINE_HEIGHT*1), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+            cv2.putText(frame, f'Frame: {i}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*0), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+
+            vid_frames.append(frame)
+            
+    return vid_frames, h, w
+
+def draw_fish(frame, left, right, top, bottom, color, fish_id, fish_len, LINE_HEIGHT=18, anno_align="left"):
+    cv2.rectangle(frame, (left, top), (right, bottom), color, BOX_THICKNESS)
+    
+    if anno_align == "left":
+        anno_align = left
+    else:
+        anno_align = right
+    cv2.putText(frame, fish_id, (anno_align, top), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, color, FONT_THICKNESS, cv2.LINE_AA, False)
+    cv2.putText(frame, fish_len, (anno_align, bottom+int(LINE_HEIGHT/2)), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, color, FONT_THICKNESS, cv2.LINE_AA, False)
+    
+def add_axis(img, image_meter_width=None, image_meter_height=None):
+    h, w, c = img.shape
+    
+    # add black border around image
+    bordersize_t = 25
+    bordersize_l = 45
+    img = cv2.copyMakeBorder(
+        img,
+        bottom=bordersize_t,
+        top=0,
+        left=bordersize_l,
+        right=25, # this helps with text getting cut off
+        borderType=cv2.BORDER_CONSTANT,
+        value=BLACK
+    )
+    
+    # add axis
+    axis_thickness = 1
+    img = cv2.line(img, (bordersize_l, h+axis_thickness//2), (w+bordersize_l, h+axis_thickness//2), WHITE, axis_thickness) # x
+    img = cv2.line(img, (bordersize_l-axis_thickness//2, 0), (bordersize_l-axis_thickness//2, h), WHITE, axis_thickness) # y
+    
+    # dist between ticks in meters
+    x_inc = 100
+    if image_meter_width and image_meter_width > 0:
+        x_inc = w / image_meter_width / 2 # 0.5m ticks
+        if image_meter_width > 4:
+            x_inc *= 2 # 1m ticks
+        if image_meter_width > 8:
+            x_inc *= 2 # 2m ticks
+            
+    # dist between ticks in meters
+    y_inc = 100
+    if image_meter_height and image_meter_height > 0:
+        y_inc = h / image_meter_height / 2 # 0.5m ticks
+        if image_meter_height > 4:
+            y_inc *= 2 # 1m ticks
+        if image_meter_height > 8:
+            y_inc *= 2 # 2m ticks
+        if image_meter_height > 12:
+            y_inc *= 3/2 # 3m ticks
+            
+    # tick mark labels
+    def x_label(x):
+        if image_meter_width and image_meter_width > 0:
+            if x_inc < w / image_meter_width: # fractional ticks
+                return "%.1fm" % (x / w * image_meter_width)
+            return "%.0fm" % (x / w * image_meter_width)
+        return str(x) # pixels
+    def y_label(y):
+        if image_meter_height and image_meter_height > 0:
+            if y_inc < y / image_meter_height: # fractional ticks
+                return "%.1fm" % (y / h * image_meter_height)
+            return "%.0fm" % (y / h * image_meter_height)
+        return str(y) # pixels
+
+    # add ticks
+    ticksize = 5
+    x = 0
+    while x < w:
+        img = cv2.line(img, (int(bordersize_l+x), h+axis_thickness//2), (int(bordersize_l+x), h+axis_thickness//2+ticksize), WHITE, axis_thickness)
+        cv2.putText(img, x_label(x), (int(bordersize_l+x), h+axis_thickness//2+LINE_HEIGHT), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE*3/4, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+        x += x_inc
+    y = 0
+    while y < h:
+        img = cv2.line(img, (bordersize_l-axis_thickness//2, int(h-y)), (bordersize_l-axis_thickness//2-ticksize, int(h-y)), WHITE, axis_thickness)
+        ylabel = y_label(y)
+        txt_offset = 13*len(ylabel)
+        cv2.putText(img, y_label(y), (bordersize_l-axis_thickness//2-ticksize - txt_offset, int(h-y)), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE*3/4, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
+        y += y_inc
+    
+    # resize to original dims
+    return cv2.resize(img, (w,h))