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

# 1. Load YOLOv8 model
model = YOLO("best_rockfall_model.pt")

# 2. Inference function (image + JSON)
def predict(image):
    results = model.predict(image)
    result = results[0]

    # --- Prepare annotated image ---
    annotated = result.plot()  # numpy array (BGR)
    annotated_img = Image.fromarray(annotated[..., ::-1])  # convert BGR → RGB

    # --- Prepare JSON detections ---
    detections = []
    for box in result.boxes:
        cls_id = int(box.cls[0])
        label = result.names[cls_id]
        conf = float(box.conf[0])
        xyxy = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
        detections.append({
            "label": label,
            "confidence": conf,
            "bbox": xyxy
        })

    return annotated_img, detections

# 3. Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Annotated Image"),
        gr.JSON(label="Detections")
    ],
    title="Rockfall Detection API (YOLOv8)",
    description="Upload an image to detect rockfalls. The model returns both the annotated image and raw detections (labels, confidence, bounding boxes)."
)

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