from glob import glob
from PIL import Image
from ultralytics import YOLO
from utils import draw_bbox
import gradio as gr
import numpy as np
import subprocess


with gr.Blocks() as demo:
    gr.Markdown("Detect planes demo.")

    models=["SSD", "FasterRCNN", "CenterNet", "RetinaNet", "DETR", "RTMDET", "YOLOv5", "YOLOv8"] 

    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column():
                image_input_single = gr.Image()
            image_output = gr.Image(visible = True)
        with gr.Row():
            drop = gr.Dropdown([m for m in models], label="Model selection", type ="index", value=models[0])
            image_button = gr.Button("Detect", variant = 'primary')
            with gr.Column(visible=True) as output_row:
                object_count = gr.Textbox(value = 0,label="Aircrafts detected")

    def runmodel(input_img, model_num):
        Image.fromarray(input_img).save(source:="inptest.jpg")
        print("Using model", model_name:=models[model_num])

        conf = 0.3

        if model_name in models[:-2]:
            cmd = f"python3 image_inference.py {source} inference/{model_name.lower()}_config.py --weights inference/models/{model_name.lower()}best.pth --out-dir inference/results/{model_name.lower()}_inference --pred-score-thr {conf}"
            subprocess.run(cmd, shell=True)
            im, count = draw_bbox(model_name.lower())

        if model_name == "YOLOv5":
            cmd = f"python3 yolov5/detect.py --weights inference/models/yolov5best.pt --source {source}  --save-txt --save-conf --project inference/results/yolov5_inference --name predict"
            subprocess.run(cmd, shell=True)
            im, count = draw_bbox(model_name.lower())

        if model_name == "YOLOv8":
            model = YOLO('inference/models/yolov8best.pt')
            results = model(source, imgsz=1024, conf = conf, save_txt = True, save_conf = True, save = True, project = "inference/results/yolov8_inference")
            im, count = draw_bbox(model_name.lower())

        return im, count

    image_button.click(runmodel, inputs=[image_input_single, drop], outputs=[image_output, object_count])

demo.launch()