from ultralytics import YOLO
import gradio as gr


model = YOLO("./runs/detect/train18/weights/best.pt")


def yolo_predict(image):
    """Run YOLOv8 inference and return annotated image with results"""
    results = model(image)
    print(results)
    annotated_image = results[0].plot()
    
    # Get prediction details
    boxes = results[0].boxes
    prediction_details = []
    
    for box in boxes:
        class_id = int(box.cls[0].item())
        class_name = model.names[class_id]
        confidence = round(box.conf[0].item(), 2)
        coords = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
        
        prediction_details.append({
            "class": class_name,
            "confidence": confidence,
            "bbox": coords
        })

    return annotated_image
    # return annotated_image, prediction_details


with gr.Blocks() as demo:
    gr.Markdown("# YOLOv8 Object Detection")
    gr.Markdown(
        """
        This application uses a YOLOv8m model fine-tuned specifically to detect red blood
        cells, white blood cells, and platelets in images of blood cells. This version
        was trained using the `keremberke/blood-cell-object-detection` dataset on huggingface.com. 
        """
    )

    gr.Interface(
        fn=yolo_predict,
        inputs=gr.Image(label="Input Image",type="pil"),
        outputs=[
            gr.Image(label="Detected Objects"),
            # gr.JSON(label="Detection Details")
        ],
        # title="YOLOv8 Object Detection",
        # # description="Upload an image to detect objects using YOLOv8",
        description='Select an example image below (none of which were included in model training or validation), or upload your own image. Then, click "Submit" to see the model in action.',
        examples=[
            "./bloodcell-examples/image_0.jpg",
            "./bloodcell-examples/image_1.jpg",
            "./bloodcell-examples/image_2.jpg",
            "./bloodcell-examples/image_3.jpg",
            "./bloodcell-examples/image_4.jpg",            
        ],
    )


demo.launch(
    show_error=True,
    height=900,
    width="80%",
    # width="100%",
    # share=True,
)