frontend/result_ui.py

import gradio as gr
import numpy as np
import os

from frontend.pdf_handler import make_pdf

js_update_tab_labels = """
    async () => {
        let el_list = document.getElementById("tab_labeler").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 += "#result_tabs button.svelte-kqij2n:nth-child(" + i + "):before {content: 'Result " + i + "';}"
        }
    }
"""

table_headers = ["TOTAL", "START_FRAME", "END_FRAME", "DETECTION_DROPOUT", "DIR", "R", "THETA", "L", "TRAVEL"]
info_headers = [
    "TOTAL_TIME", "DATE", "START", "END", "FRAME_RATE", "",
    "TOTAL_FISH", "UPSTREAM_FISH", "DOWNSTREAM_FISH", "NONDIRECTIONAL_FISH", "",
    "UPSTREAM_MOTION", "INTENSITY", "THRESHOLD", "WATER_TEMP", "",
]
max_tabs = 10

tabs = []
tab_parent = None
zip_out = None

def update_result(i, state, result, inference_handler):
    print("loading result")
    # If index is larger than max_tabs, only add file to zip list
    if i >= max_tabs:
        return {
            zip_out: gr.update(value=result["path_zip"])
        }

    # Check if inference is done
    not_done = state['index'] < state['total']

    annotation_avaliable = state['enable_annotation_editor'] and (result["aris_input"][i] is not None)

    if 'Generate PDF' in state['outputs']:
        print("making pdf")
        make_pdf(state['index']-1, state, result, table_headers)
        print("done pdf")

    # Check if files exist
    video_path = result["path_video"][i]
    if not os.path.exists(video_path): video_path = None

    # Send update to UI, and to inference_handler to start next file inference

    print("loaded result")
    return {
        zip_out: gr.update(value=result["path_zip"]),
        tabs[i]['tab']: gr.update(),
        tabs[i]['video']: gr.update(value=video_path, visible=video_path is not None),
        tabs[i]['metadata']: gr.update(value=result["fish_info"][i], visible=True),
        tabs[i]['table']: gr.update(value=result["fish_table"][i], visible=True),
        tabs[i]['annotation_btn']: gr.update(visible=annotation_avaliable),
        tab_parent: gr.update(selected=i),
        inference_handler: gr.update(value = str(np.random.rand()), visible=not_done)
    }


# Auto_download
def auto_download_zip(state):
    if 'Automatically download result' in state['outputs']:
        return gr.update(value=str(np.random.rand()))
    else:
        return gr.update()

def Result_Gradio(prepare_annotation, components, state):
    global tabs, tab_parent, zip_out

    # Dummy element to call inference events, this also displays the inference progress
    components['inference_handler'] = gr.Text(value=str(np.random.rand()), visible=False)

    # Dummy element to call UI events
    components['result_handler'] = gr.Text(value="LOADING", visible=False)

    # Dummy element for updating tab titles
    components['tab_labeler'] = gr.Text(value="", visible=False, elem_id="tab_labeler")
    components['tab_labeler'].change(lambda x: x, None, None, _js=js_update_tab_labels)

    components['cancel_btn'] = gr.Button("Cancel Inference", visible=False)

    # List of all UI components that will recieve outputs from the result_handler
    visual_components = []

    # Zip file output
    zip_out = gr.File(label="ZIP Output", elem_id="zip_out", interactive=False)
    visual_components.append(zip_out)
    components['zip_out'] = zip_out

    autodownloader = gr.Text(value="LOADING", visible=False)

    zip_out.change(lambda: auto_download_zip(state), None, autodownloader)
    autodownloader.change(lambda x: x, autodownloader, None, _js="""
        () => {
            zip_out = document.getElementById("zip_out")
            downloads = zip_out?.getElementsByClassName("download")
            if (downloads?.length > 0) {
                downloads[downloads.length-1].children[0].click()       
            }
        }
        """
    )

    # Create result tabs
    tabs = []
    with gr.Tabs(elem_id="result_tabs") as tab_parent:
        visual_components.append(tab_parent)

        # Create 'max_tab' tabs for showing result
        for i in range(max_tabs):
            with gr.Tab(label="", id=i, elem_id="result_tab"+str(i)) as tab:

                with gr.Row():
                    # List of clip info (date, time, number of fish, temperature, etc.) 
                    metadata_out = gr.Matrix(label="Info", interactive=False, headers=["Field", "Value"], datatype="markdown", visible=False, elem_id="marking_json")
                    
                    # Annotated video
                    video_out = gr.Video(label='Annotated Video', interactive=False, visible=False)


                # Table of found fish
                table_out = gr.Matrix(label='Indentified Fish', headers=table_headers, interactive=False, visible=False)

                # Button for opening result in annotation editor
                annotation_btn = gr.Button("Edit Annotation", visible=False)
                annotation_btn.click(prepare_annotation, annotation_btn, [components['annotation_progress'], components['master_tabs']], _js="() => " + str(i))
                
                # Add components to tab dict for easy access later on
                tabs.append({
                    'tab': tab,
                    'metadata': metadata_out,
                    'video': video_out,
                    'table': table_out,
                    'annotation_btn': annotation_btn
                })

                # Add all components to list of visualization outputs
                visual_components.extend([tab, metadata_out, video_out, table_out, annotation_btn])

    components['result_tabs'] = tab_parent

    return visual_components


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

    # Create fish table
    table = []
    for fish in metadata["FISH"]:
        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)

    # Create info table
    info = []
    for field in info_headers:
        field_name = "**" + field + "**"
        if field in metadata:
            info.append([field_name, str(metadata[field])])
        else:
            info.append([field_name, ""])
    if 'hyperparameters' in metadata:
        for param_name in metadata['hyperparameters']:
            info.append(['**' + param_name + '**', str(metadata['hyperparameters'][param_name])])
    
    return table, info