import gradio as gr
from predict import ONNXInference

PRED = []

def detect(files):
    model = ONNXInference(
        model_path="/home/neo/Downloads/torchFlow/models/torchFlow-ckpt.onnx",
        files=files,
        save_image=False,
        save_path="/home/neo/Downloads/torchFlow/"
    )
    res = model.run()
    img_id = res["IMG_ID"]
    pred_lab = res["PRED_LAB"],
    pred_ct = res["PRED_CT"],
    geo_tag_url =  res["GEO_TAG_URL"]
    PRED.append(pred_ct)
    return f"Predicted"

with gr.Blocks() as demo:
    with gr.Row():
        output=gr.Image()
    with gr.Row():
        btn = gr.UploadButton(
            label="Upload Image",
            file_types = ['.jpg','.jpeg'],
            file_count = "multiple")
        btn.upload(fn=detect, inputs=btn)
        with gr.Column(scale=1, min_width=600):
            gr.Markdown(f"Output here")
            if PRED is not None:
                gr.Markdown(f"Predicted: {PRED}")

demo.launch()