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

# Load the YOLO model - YOLOv11m for pothole, road damage, and garbage detection
try:
    model = YOLO("Vision Classification.pt")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

def predict(image, conf_threshold):
    if image is None or model is None:
        return None, "Model not loaded or invalid image."
    
    # Run inference (imgsz=768 based on model card)
    results = model(image, imgsz=768, conf=conf_threshold)
    
    # YOLO returns a list of Results objects
    result = results[0]
    
    # Plotting the detections on the image
    # plot() returns a BGR numpy array
    annotated_image = result.plot()
    
    # Convert BGR to RGB for Gradio Display
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    # Detection overview text
    boxes = result.boxes
    class_names = result.names
    
    if len(boxes) == 0:
        detection_summary = "No civic issues detected in this image."
    else:
        # Count detections
        detection_counts = {}
        for box in boxes:
            cls_id = int(box.cls)
            cls_name = class_names[cls_id]
            detection_counts[cls_name] = detection_counts.get(cls_name, 0) + 1
            
        summary_lines = ["**Detections:**"]
        for cls_name, count in detection_counts.items():
            summary_lines.append(f"- {count} {cls_name}(s)")
            
        detection_summary = "\n".join(summary_lines)
    
    return Image.fromarray(annotated_image_rgb), detection_summary

# Gradio Interface
with gr.Blocks(title="PotholeNet-YOLO11m-v1 🛑") as interface:
    gr.Markdown("# 🛑 PotholeNet-YOLO11m-v1")
    gr.Markdown("**Aamchi City AI Civic System** — Real-time pothole, road damage, and garbage detection for Indian urban roads.")
    gr.Markdown("Upload an image of a road to detect infrastructure issues. The model was trained on 23,000+ street-level images.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Street Image")
            conf_slider = gr.Slider(minimum=0.01, maximum=1.0, value=0.25, step=0.01, label="Confidence Threshold")
            submit_btn = gr.Button("Detect Civic Issues", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(type="pil", label="Detection Results")
            detection_text = gr.Markdown(label="Detection Summary")
            
    submit_btn.click(
        fn=predict,
        inputs=[input_image, conf_slider],
        outputs=[output_image, detection_text]
    )
    
    gr.Markdown("### Intended Use")
    gr.Markdown("Real-time pothole detection, Automated civic issue reporting, Infrastructure health monitoring.")
    gr.Markdown("**Developer:** Vansh Momaya")

if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)