oskarastrom's picture
Compatability with backend
128e4f0
raw
history blame contribute delete
No virus
5.46 kB
import gradio as gr
from frontend.custom_file_reader import File
from backend.InferenceConfig import InferenceConfig, TrackerType
import os
models = {
'master': 'models/v5m_896_300best.pt',
# 'elwha': 'models/YsEE20.pt',
# 'elwha+kenai_val': 'models/YsEKvE20.pt',
'elwha': 'models/YsEKtE20.pt',
}
def Upload_Gradio(gradio_components):
with gr.Tabs():
# Tab - uploading aris files for inference
with gr.Tab("Infer ARIS"):
gr.HTML("<p align='center' style='font-size: large;font-style: italic;'>Submit an .aris file to analyze result.</p>")
default_settings = InferenceConfig()
hyperparams = []
with gr.Accordion("Advanced Settings", open=False):
default_model = default_settings.find_model(models)
hyperparams.append(gr.Dropdown(label="Model", value=default_model, choices=list(models.keys())))
gr.Markdown("Detection Parameters")
with gr.Row():
hyperparams.append(gr.Slider(0, 1, value=default_settings.conf_thresh, label="Confidence Threshold", info="Confidence cutoff for detection boxes"))
hyperparams.append(gr.Slider(0, 1, value=default_settings.nms_iou, label="NMS IoU", info="IoU threshold for non-max suppression"))
gr.Markdown("Tracking Parameters")
with gr.Row():
hyperparams.append(gr.Slider(0, 100, value=default_settings.min_hits, label="Min Hits", info="Minimum number of frames a fish has to appear in to count"))
hyperparams.append(gr.Slider(0, 100, value=default_settings.max_age, label="Max Age", info="Max age of occlusion before track is split"))
default_tracker = TrackerType.toString(default_settings.associative_tracker)
tracker = gr.Dropdown(["None", "Confidence Boost", "ByteTrack"], value=default_tracker, label="Associative Tracking")
hyperparams.append(tracker)
with gr.Row(visible=default_tracker=="Confidence Boost") as track_row:
hyperparams.append(gr.Slider(0, 5, value=default_settings.boost_power, label="Boost Power", info="Scalar multiplier for the boost amount"))
hyperparams.append(gr.Slider(0, 1, value=default_settings.boost_decay, label="Boost Decay", info="Exponential decay parameter for boost based on frame time difference"))
tracker.change(lambda x: gr.update(visible=(x=="Confidence Boost")), tracker, track_row)
with gr.Row(visible=default_tracker=="ByteTrack") as track_row:
hyperparams.append(gr.Slider(0, 1, value=default_settings.byte_low_conf, label="Low Conf Threshold", info="Confidence threshold for the low detection group"))
hyperparams.append(gr.Slider(0, 1, value=default_settings.byte_high_conf, label="High Conf Threshold", info="Confidence threshold for the high detection group"))
tracker.change(lambda x: gr.update(visible=(x=="ByteTrack")), tracker, track_row)
gr.Markdown("Other")
with gr.Row():
hyperparams.append(gr.Slider(0, 3, value=default_settings.min_length, label="Min Length", info="Minimum length of fish (meters) in order for it to count"))
hyperparams.append(gr.Slider(0, 3, value=default_settings.max_length, label="Max Length", info="Maximum length of fish (meters) in order for it to count. (disable at 0)"))
hyperparams.append(gr.Slider(0, 10, value=default_settings.min_travel, label="Min Travel", info="Minimum travel distance of track (meters) in order for it to count"))
gradio_components['hyperparams'] = hyperparams
with gr.Row():
hyperparams.append(gr.CheckboxGroup([("Generate Annotated Video"), ("Generate Manual Marking"), ("Generate PDF"), ("Automatically download result")], label="Output settings", interactive=True, value=["Generate Annotated Video"]))
#Input field for aris submission
gradio_components['input'] = File(file_types=[".aris", ".ddf"], type="binary", label="ARIS Input", file_count="multiple")
example_name = "static/example.aris"
gradio_components['examples'] = gr.Examples(examples=[[example_name]], inputs=gradio_components['input'])
gradio_components['inference_btn'] = gr.Button("Run")
# Tab - uploading old result files to review
with gr.Tab("Open Result"):
gr.HTML("""
<p align='center' style='font-size: large;font-style: italic;'>Submit an old zip file of results to visualize.</p>
<p align='center' style='font-size: large;font-style: italic;'>If you want to edit annotations, also submit an aris file.</p>
""")
# Input for .zip result file
gradio_components['result_input'] = File(file_types=[".zip"], type="binary", label="Upload result file", file_count="multiple")
# Optional input for aris file to help with annotation editing
gradio_components['result_aris_input'] = File(file_types=[".aris", ".ddf"], type="binary", label="Upload aris file (optional)", file_count="multiple")
# Button for initializing review
gradio_components['open_result_btn'] = gr.Button("View Result")