import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
from PIL import Image

MODEL_REPOS = {
    "Footprint YOLO": "risashinoda/footprint_yolo",     
    "Feces YOLO": "risashinoda/feces_yolo",     
    "Egg YOLO": "risashinoda/egg_yolo",       
    "Bone YOLO": "risashinoda/bone_yolo",      
    "Feather YOLO": "risashinoda/feather_yolo"
}

_loaded = {}

def _load(model_key, weights_name="last.pt"):
    if model_key not in _loaded:
        repo_id = MODEL_REPOS[model_key]
        w = hf_hub_download(repo_id=repo_id, filename=weights_name)
        _loaded[model_key] = YOLO(w)
    return _loaded[model_key]

def infer(image, model_key, conf_thres=None, iou_nms=None, draw_labels=None):
    conf_thres = 0.25 if conf_thres is None else float(conf_thres)
    iou_nms    = 0.70 if iou_nms is None else float(iou_nms)
    draw_labels = True if draw_labels is None else bool(draw_labels)

    model = _load(model_key)
    results = model.predict(image, conf=conf_thres, iou=iou_nms)
    r = results[0]

    plotted = r.plot()
    img_out = Image.fromarray(plotted[..., ::-1])

    if not draw_labels:
        import numpy as np, cv2
        img = results[0].orig_img.copy()
        if img.ndim == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        for box in r.boxes.xyxy.tolist():
            x1, y1, x2, y2 = map(int, box)
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 150, 0), 3)
        img_out = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

    return img_out

with gr.Blocks() as demo:
    gr.Markdown("# Multi-YOLO Demo (BBox only)")
    gr.Markdown("## 🔎 **Select a model first**\nChoose one below, then upload an image.")

    with gr.Row():
        img_in = gr.Image(type="pil", label="Upload an image")

        with gr.Column():
            # ここは Radio にする
            model_dd = gr.Radio(
                choices=list(MODEL_REPOS.keys()),
                value=list(MODEL_REPOS.keys())[0],
                label="Select Model",
                interactive=True
            )

            conf = gr.Slider(0.05, 1.0, value=0.25, step=0.01, label="Confidence (default 0.25)")
            iou = gr.Slider(0.1, 0.95, value=0.70, step=0.01, label="NMS IoU (default 0.70)")
            draw_labels = gr.Checkbox(value=True, label="Draw labels text (off = boxes only)")
            run_btn = gr.Button("Run")

    img_out = gr.Image(type="pil", label="Detections (boxes only)")

    run_btn.click(fn=infer, inputs=[img_in, model_dd, conf, iou, draw_labels], outputs=[img_out])

if __name__ == "__main__":
    demo.launch()