import gradio as gr
import cv2
import torch
import numpy as np

# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# Function to run inference on an image and count objects
def run_inference(image):
    # Convert the image from PIL format to a format compatible with OpenCV
    image = np.array(image)

    # Run YOLOv5 inference
    results = model(image)

    # Extract detection results
    detections = results.pandas().xyxy[0]

    # Count objects by category
    object_counts = detections['name'].value_counts()

    # Create a formatted string to show object counts
    count_text = "\n".join([f"{obj}: {count}" for obj, count in object_counts.items()])

    # Convert the annotated image from BGR to RGB for display
    annotated_image = results.render()[0]
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

    return annotated_image, count_text

# Create the Gradio interface
interface = gr.Interface(
    fn=run_inference,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil"),
        gr.Textbox(label="Object Counts", lines=5, interactive=False)  # Display counts as text
    ],
    title="YOLOv5 Object Detection with Counts",
    description="Upload an image to run YOLOv5 object detection, see the annotated results, and view the count of detected objects by category."
)

# Launch the app
interface.launch()