import gradio as gr
import numpy as np
import json
from zipfile import ZipFile
import os
from backend.dataloader import create_dataloader_aris
from backend.aws_handler import ping_server
from backend.predict import predict_task
from backend.uploader import save_data_to_dir, create_data_dir, save_data
from backend.InferenceConfig import InferenceConfig
from frontend.upload_ui import Upload_Gradio, models
from frontend.result_ui import Result_Gradio, update_result, create_metadata_table, table_headers, info_headers
from frontend.annotation_handler import load_annotation, prepare_annotation, js_store_frame_info, annotation_css
from frontend.state_handler import reset_state
#Initialize State & Result
state = {
'version': "Experimental 1.0",
'files': [],
'index': 1,
'total': 1,
'annotation_index': -1,
'frame_index': 0,
'outputs': [],
'config': None,
'enable_annotation_editor': False,
}
result = {}
components = {}
# -------------------------------------------- ----- UPLOAD ARIS FILE ------------------------------------------------------
# Called when an Aris file is uploaded for inference - calls infer_next
def on_aris_input(
file_list,
model_id,
conf_thresh, iou_thresh,
min_hits, max_age,
associative_tracker, boost_power, boost_decay, byte_low_conf, byte_high_conf,
min_length, max_length, min_travel,
output_formats
):
if isinstance(file_list, tuple):
file_list = [file_list]
print(output_formats)
# Reset Result
reset_state(result, state)
state['files'] = file_list
state['total'] = len(file_list)
state['outputs'] = output_formats
state['config'] = InferenceConfig(
weights = models[model_id] if model_id in models else models['master'],
conf_thresh = conf_thresh,
nms_iou = iou_thresh,
min_hits = min_hits,
max_age = max_age,
min_length = min_length,
max_length = max_length,
min_travel = min_travel,
)
# Enable tracker if specified
if (associative_tracker == "Confidence Boost"):
state['config'].enable_conf_boost(boost_power, boost_decay)
elif (associative_tracker == "ByteTrack"):
state['config'].enable_byte_track(byte_low_conf, byte_high_conf)
else:
state['config'].enable_sort_track()
print(" ")
print("Inference with:")
print(state['config'].to_dict())
print(" ")
# Update loading_space to start inference on first file
return {
inference_handler: gr.update(value = str(np.random.rand()), visible=True),
components['cancel_btn']: gr.update(visible=True),
master_tabs: gr.update(selected=1)
}
# Iterative function that performs inference on the next file in line
def infer_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...")
# Save file and create a new directory for result
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("Directory: ", dir_name)
print("Aris input: ", file_path)
print(" ")
# Check that the file was valid
if not valid:
return {
result_handler: gr.update(),
inference_handler: gr.update()
}
# Send uploaded file to AWS
ping_server(file_name, state)
#upload_file(file_path, "fishcounting", "webapp_uploads/files/" + file_name)
#crop_clip(file_path, 65)
# Do inference
json_result, json_filepath, zip_filepath, video_filepath, marking_filepath = predict_task(
file_path,
config = state['config'],
output_formats = state['outputs'],
gradio_progress = set_progress
)
# prepare dummy dataloader for visualizations
_, dataset = create_dataloader_aris(file_path, num_frames_bg_subtract=0)
# Store result for that file
result['json_result'].append(json_result)
result['aris_input'].append(file_path)
result['datasets'].append(dataset)
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(json_result, table_headers, info_headers)
result["fish_table"].append(fish_table)
result["fish_info"].append(fish_info)
# Increase file index
state['index'] += 1
# Send of update to result_handler to show new result
# Leave inference_handler update blank to avoid starting next inference until result is updated
return {
result_handler: gr.update(value = str(np.random.rand())),
tab_labeler: gr.update(value = str(state['index'])),
inference_handler: gr.update()
}
# Cancel inference
def cancel_inference():
return {
master_tabs: gr.update(selected=0),
inference_handler: gr.update(visible=False),
components['cancel_btn']: gr.update(visible=False)
}
# Show result
def on_result_ready():
# Update result tab for last file
i = state["index"] - 1
return update_result(i, state, result, inference_handler)
# ------------------------------------------------- UPLOAD RESULT FILE -----------------------------------------------------
# Called when result file is uploaded for review
def on_result_upload():
return {
master_tabs: gr.update(selected=1),
result_uploader: gr.update(value=str(np.random.rand()))
}
# Called when result upload is finished processing
def on_result_upload_finish(zip_list, aris_list):
if (zip_list == None):
zip_list = [("static/example/example_result.zip", None)]
aris_path = "static/example/input_file.aris"
aris_list = [(aris_path, bytearray(open(aris_path, 'rb').read()))]
reset_state(result, state)
state['outputs'] = ["Generate Annotated Video", "Generate Manual Marking", "Generate PDF"]
component_updates = {
tab_labeler: gr.update(value = len(zip_list))
}
for i in range(len(zip_list)):
# Create dir to unzip files
dir_name = create_data_dir(str(i))
# Check aris input
if (aris_list):
aris_info = aris_list[i]
file_name = aris_info[0].split("/")[-1]
bytes = aris_info[1]
valid, input_path, dir_name = save_data_to_dir(bytes, file_name, dir_name)
_, dataset = create_dataloader_aris(input_path, num_frames_bg_subtract=0)
else:
input_path = None
dataset = None
# Unzip result
zip_info = zip_list[i]
zip_name = zip_info[0]
print(zip_name)
with ZipFile(zip_name) as zip_file:
ZipFile.extractall(zip_file, path=dir_name)
unzipped = os.listdir(dir_name)
print(unzipped)
for file in unzipped:
if (file.endswith("_results.mp4")):
result["path_video"].append(os.path.join(dir_name, file))
elif (file.endswith("_results.json")):
result["path_json"].append(os.path.join(dir_name, file))
elif (file.endswith("_marking.txt")):
result["path_marking"].append(os.path.join(dir_name, file))
result["aris_input"].append(input_path)
result["datasets"].append(dataset)
with open(result['path_json'][-1]) as f:
json_result = json.load(f)
result['json_result'].append(json_result)
fish_table, fish_info = create_metadata_table(json_result, table_headers, info_headers)
result["fish_table"].append(fish_table)
result["fish_info"].append(fish_info)
update = update_result(i, state, result, inference_handler)
for key in update.keys():
component_updates[key] = update[key]
component_updates.pop(inference_handler)
return component_updates
# ------------------------------------------------- ANNOTATION EDITOR -----------------------------------------------------
def on_annotation_open(result_index):
return prepare_annotation(state, result, result_index)
def annotate_next(_, progress=gr.Progress()):
return load_annotation(state, result, progress)
# -------------------------------------------------- GRADIO ARCHITECTURE ----------------------------------------------------
with gr.Blocks() as demo:
with gr.Blocks():
# Title of page + style
gr.HTML(
"""
<h1 align="center" style="font-size:xxx-large">Caltech Fisheye - Experimental</h1>
<style>
/* Disable header of metadata list in result */
#marking_json thead {
display: none !important;
}
/* Color of selected tab */
.selected.svelte-kqij2n {
background: linear-gradient(180deg, #66eecb47, transparent);
}
""" + annotation_css + """
</style>
<style id="tab_style"></style>
"""
)
with gr.Tabs() as master_tabs:
components['master_tabs'] = master_tabs
# Master Tab for uploading aris or result files
with gr.Tab("Upload", id=0):
# Draw Gradio components related to the upload ui
Upload_Gradio(components)
# Master Tab for result visualization
with gr.Tab("Result", id=1):
# Define annotation progress bar for event listeres, but unrender since it will be displayed later on
result_uploader = gr.HTML("", visible=False)
components['result_uploader'] = result_uploader
annotation_progress = gr.HTML("", visible=False).unrender()
components['annotation_progress'] = annotation_progress
# Draw the gradio components related to visualzing result
visualization_components = Result_Gradio(on_annotation_open, components, state)
# Master Tab for annotation editing
if state['enable_annotation_editor']:
with gr.Tab("Annotation Editor", id=2):
# Draw the annotation loading bar here
annotation_progress.render()
# Add annotation editor component
annotation_editor = gr.HTML("", visible=False)
# Event listener for batch loading of annotation frames
annotation_progress.change(
annotate_next,
annotation_progress,
[annotation_editor, annotation_progress],
_js=js_store_frame_info
)
# Event listener for running javascript defined in 'annotation_editor.js'
# show_annotation
with open('gradio_scripts/annotation_editor.js', 'r') as f:
annotation_editor.change(lambda x: gr.update(), None, annotation_editor, _js=f.read())
# 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>
"""
)
# Extract important components for ease of code
input = components['input']
inference_handler = components['inference_handler']
result_handler = components['result_handler']
tab_labeler = components['tab_labeler']
inference_comps = [inference_handler, master_tabs, components['cancel_btn']]
# When a file is uploaded to the input, tell the inference_handler to start inference
input.upload(on_aris_input, [input] + components['hyperparams'], inference_comps)
components['inference_btn'].click(on_aris_input, [input] + components['hyperparams'], inference_comps)
# 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_event = inference_handler.change(infer_next, None, [inference_handler, result_handler, tab_labeler])
# Send UI changes based on the new results to the UI_components, and tell the inference_handler to start next inference
result_handler.change(on_result_ready, None, visualization_components + [inference_handler])
# Cancel and skip buttons
components['cancel_btn'].click(cancel_inference, None, inference_comps, cancels=[inference_event])
# Button to load a previous result and view visualization
components['open_result_btn'].click(on_result_upload, None, [result_uploader, master_tabs])
components['result_uploader'].change(
on_result_upload_finish,
[components['result_input'], components['result_aris_input']],
visualization_components + [tab_labeler]
)
demo.queue().launch()