app.py

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
import cv2

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 = {}


out = cv2.VideoWriter("static/test_video.mp4", cv2.VideoWriter_fourcc(*'mp4v'), 10, [ 100, 100 ] )
for i in range(20):
    frame = (np.random.rand(100, 100, 3)*255).astype('uint8')
    out.write(frame)
out.release()



# 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']

    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)
        
        video_out = gr.Video(value="static/test_video.mp4", label='Annotated Video', interactive=False)

        # 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()