app.py

import gradio as gr
from ultralytics import YOLO
from PIL import Image
import tempfile
import os

def detect(weights_file, image_file, conf_threshold):
    if weights_file is None or image_file is None:
        return None, "Please upload both a .pt weights file and an image."

    # Save uploaded weights to a temp path
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_w:
        tmp_w.write(weights_file.read())
        weights_path = tmp_w.name

    # Load model from uploaded weights
    model = YOLO(weights_path)

    # Load input image as PIL
    if isinstance(image_file, str):
        img = Image.open(image_file).convert("RGB")
    else:
        img = Image.open(image_file).convert("RGB")

    # Run prediction; return annotated image
    results = model.predict(
        source=img,
        conf=conf_threshold,
        imgsz=640,
        verbose=False,
    )
    r = results[0]
    annotated = r.plot()        # numpy array BGR
    annotated_pil = Image.fromarray(annotated[:, :, ::-1])  # BGR -> RGB

    return (img, annotated_pil), f"Detections done with {os.path.basename(weights_file.name)}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# YOLOv8 Viewer\nUpload a YOLO `.pt` weights file and an image. The app will show original and detections side by side.")

    with gr.Row():
        weights_input = gr.File(
            label="YOLO weights (.pt)",
            file_types=[".pt"],
            type="binary",
        )
        conf_slider = gr.Slider(
            minimum=0.1,
            maximum=0.9,
            value=0.25,
            step=0.05,
            label="Confidence threshold",
        )

    image_input = gr.Image(type="filepath", label="Input image")

    gallery = gr.Gallery(
        label="Original (left) vs Detections (right)",
        columns=2,
        height=512,
    )
    status = gr.Textbox(label="Status / Info", interactive=False)

    run_btn = gr.Button("Run detection")

    run_btn.click(
        fn=detect,
        inputs=[weights_input, image_input, conf_slider],
        outputs=[gallery, status],
    )

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