import gradio as gr
import torch
from ultralytics import YOLO
import cv2

# Load the YOLOv8 model
model = YOLO("best.pt")

# Inference function
def predict(image):
    # Run prediction
    results = model(image)

    # Annotated image
    annotated_img = results[0].plot()

    # Prepare detection summary
    detections = results[0].boxes
    output_text = ""

    if detections is not None and len(detections.cls) > 0:
        output_text += "Prediction Summary:\n\n"
        for i, box in enumerate(detections):
            cls_id = int(box.cls.item())
            conf = float(box.conf.item())
            label = model.names[cls_id]
            health_status = "Diseased" if label.lower() != "healthy" else "Healthy"

            output_text += f"Status: {health_status}\n"
            output_text += f"Disease: {label}\n"
            output_text += f"Confidence: {conf:.2f}\n\n"
    else:
        output_text = "No disease detected. The cow appears to be healthy."

    return annotated_img, output_text

# Gradio Interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Annotated Image"),
        gr.Textbox(label="Prediction Details")
    ],
    title="CowSense - Livestock Disease Detection",
    description="Upload an image of a cow to detect health status and disease type using a trained YOLOv8 model."
)

iface.launch()