import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

# Load your local model
model = YOLO('best.pt')

def get_confidence_color(confidence):
    if confidence > 0.7:
        return (0, 255, 0)  # Green for high confidence
    elif confidence > 0.4:
        return (0, 255, 255)  # Yellow for medium confidence
    else:
        return (0, 0, 255)  # Red for low confidence

def draw_circle_with_count(image, box, count, confidence):
    x1, y1, x2, y2 = map(int, box)
    center_x = int((x1 + x2) / 2)
    center_y = int((y1 + y2) / 2)
    
    color = get_confidence_color(confidence)
    radius = 20
    cv2.circle(image, (center_x, center_y), radius, color, 2)
    
    text = str(count)
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.7
    thickness = 2
    (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
    
    text_x = center_x - text_width // 2
    text_y = center_y + text_height // 2
    
    cv2.putText(image, text, (text_x, text_y), font, font_scale, color, thickness)
    
    return image

def process_image(image, imgsz, conf, iou):
    if image is None:
        return None, "No image provided"
    
    # Convert to PIL Image if it's not already
    if isinstance(image, np.ndarray):
        image_pil = Image.fromarray(image)
    else:
        image_pil = image
    
    # Run inference with parameters
    results = model(
        image_pil,
        imgsz=int(imgsz),
        conf=conf,
        iou=iou,
    )
    
    # Convert image to numpy array for drawing
    image_np = np.array(image_pil)
    
    # Count objects and draw circles
    detection_counts = {}
    total_count = 0
    
    for result in results:
        for box in result.boxes:
            total_count += 1
            class_id = int(box.cls[0])
            class_name = model.names[class_id]
            confidence = float(box.conf[0])
            bbox = box.xyxy[0].tolist()
            
            image_np = draw_circle_with_count(image_np, bbox, total_count, confidence)
            
            if class_name not in detection_counts:
                detection_counts[class_name] = {
                    'count': 0,
                    'low_conf': 0,
                    'med_conf': 0,
                    'high_conf': 0
                }
            
            detection_counts[class_name]['count'] += 1
            if confidence > 0.7:
                detection_counts[class_name]['high_conf'] += 1
            elif confidence > 0.4:
                detection_counts[class_name]['med_conf'] += 1
            else:
                detection_counts[class_name]['low_conf'] += 1
    
    count_text = "Total Objects Detected: " + str(total_count) + "\n\n"
    count_text += "Detections by class:\n"
    for class_name, stats in detection_counts.items():
        count_text += f"\n{class_name}:\n"
        count_text += f"  Total: {stats['count']}\n"
        count_text += f"  High confidence (>70%): {stats['high_conf']} 🟢\n"
        count_text += f"  Medium confidence (40-70%): {stats['med_conf']} 🟡\n"
        count_text += f"  Low confidence (<40%): {stats['low_conf']} 🔴\n"
    
    return image_np, count_text

# Create the Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("# Image Object Detection with Confidence Indicators")
    
    with gr.Row():
        # Left column with input controls
        with gr.Column(scale=1):
            # Image input
            image_input = gr.Image(
                label="Upload Image",
                type="numpy",
                height=300
            )
            
            # Parameters
            with gr.Group():
                imgsz_dropdown = gr.Dropdown(
                    choices=["320", "640"],
                    value="640",
                    label="Image Size (px)"
                )
                
                conf_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.25,
                    step=0.05,
                    label="Confidence Threshold"
                )
                
                iou_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.45,
                    step=0.05,
                    label="IoU Threshold"
                )
        
        # Right column with results
        with gr.Column(scale=2):
            # Processed image output
            processed_output = gr.Image(
                label="Processed Image",
                height=400
            )
            
            # Detection results
            text_output = gr.Textbox(
                label="Detection Results",
                lines=10
            )
    
    # Process button
    process_btn = gr.Button("Process Image")
    
    # Status message
    status_text = gr.Textbox(label="Status", lines=1)

    def process_uploaded_image(image, imgsz, conf, iou):
        """Process the uploaded image"""
        if image is None:
            return None, None, "No image to process"
        
        processed_image, detection_text = process_image(
            image,
            imgsz,
            conf,
            iou
        )
        
        return processed_image, detection_text, "Image processed successfully"

    # Connect the process button
    process_btn.click(
        fn=process_uploaded_image,
        inputs=[
            image_input,
            imgsz_dropdown,
            conf_slider,
            iou_slider
        ],
        outputs=[
            processed_output,
            text_output,
            status_text
        ]
    )

# Launch the interface
if __name__ == "__main__":
    iface.queue()
    iface.launch()