import gradio as gr
import cv2
import numpy as np
from collections import Counter
from ultralytics import YOLO

# Load YOLOv10 model
model_path = "best.pt" 
model = YOLO(model_path)

# Define the predict function
def predict(image):
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    result = model.predict(source=image_rgb, imgsz=640, conf=0.25)

    annotated_img = result[0].plot()

    detections = result[0].boxes.data
    class_names = [model.names[int(cls)] for cls in detections[:, 5]]
    count = Counter(class_names)

    detection_str = ', '.join([f"{name}: {count}" for name, count in count.items()])
    annotated_img = annotated_img[:, :, ::-1]

    return annotated_img, detection_str

# Create Gradio interface
app = gr.Interface(
    predict,
    inputs=gr.Image(type="numpy", label="Upload an image"),
    outputs=[gr.Image(type="numpy", label="Annotated Image"), gr.Textbox(label="Detection Count")],
    title="Blood Cell Count using YOLO V10",
    description="Upload an image,then YOLO V10 model will detect and annotate blood cells."
)

# Launch the app
if __name__ == "__main__":
    app.launch(share=True, server_port=8080, debug=True)