# Created by yarramsettinaresh GORAKA DIGITAL PRIVATE LIMITED at 24/10/24
import gradio as gr
import cv2
import time
from ultralytics import YOLO

# Load your YOLO model (adjust model path or type as needed)
model_path = "model_- 11 october 2024 11_07.pt"
model = YOLO(model_path)


def ultralytics_predict(model, frame):
    confidence_threshold = 0.2
    start_time = time.time()
    results = model(frame)  # Perform inference on the frame
    end_time = time.time()

    duration = end_time - start_time
    print(f"Prediction duration: {duration:.4f} seconds")
    duration_str = f"{duration:.4f} S"

    object_count = {}  # Dictionary to store counts of detected objects

    for detection in results[0].boxes:  # Iterate through detections
        conf = float(detection.conf[0])  # Confidence score
        if conf > confidence_threshold:
            conf, pos, text, color = ultralytics(detection, duration_str)
            cv2.rectangle(frame, pos[0], pos[1], color, 2)
            cv2.putText(frame, text, (pos[0][0], pos[0][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            # Update object count
            class_id = int(detection.cls[0])
            class_name = model.names[class_id]
            if class_name not in object_count:
                object_count[class_name] = dict(count=0, objects=[])
            object_mapp = object_count[class_name]
            object_mapp["count"] = object_mapp.get("count", 0) + 1
            object_mapp["objects"].append(dict(conf=conf, pos=pos, text=text, color=color))

    return frame  # Return the count of detected objects


def ultralytics(detection, duration):
    COLOUR_MAP = {
        0: (0, 0, 255),  # Red in BGR format
        1: (0, 255, 0)  # Green in BGR format
    }

    conf = float(detection.conf[0])  # Confidence score
    class_id = int(detection.cls[0])  # Class ID
    name = model.names[class_id]  # Get class name
    xmin, ymin, xmax, ymax = map(int, detection.xyxy[0])  # Bounding box coordinates
    color = COLOUR_MAP.get(class_id, (255, 255, 255))  # Default to white if not found

    # Draw bounding box and label on the frame
    pos = (xmin, ymin), (xmax, ymax)
    text = f"{name} {round(conf, 2)} :{duration}"

    return conf, pos, text, color


def process_frame(frame):
    object_count = ultralytics_predict(model, frame)
    return frame, object_count  # Return frame and object count


def detect_image(image):
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR format for OpenCV
    result_frame, object_count = process_frame(image)
    result_frame = cv2.cvtColor(result_frame, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio
    return result_frame, object_count  # Return both the frame and the count


# Create Gradio Interface
gr.Interface(
    fn=detect_image,
    inputs=gr.Image(type="numpy"),  # Updated input format
    outputs=[
        gr.Image(type="numpy"),  # Image output
        gr.JSON(),  # Object count output as JSON
    ],
    title="YOLO Object Detection",
    description="Upload an image to detect objects using YOLO model."
).launch()