app.py

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("model.pt")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

def predict(image, conf_threshold):
    try:
        if image is None or model is None:
            return None, "Model not loaded or invalid image."
        
        # Run inference
        results = model(image, imgsz=768, conf=conf_threshold)
        result = results[0]
        
        # Plotting the detections on the image returns a BGR numpy array
        annotated_image = result.plot()
        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 safely
            detection_counts = {}
            for box in boxes:
                # box.cls is usually a tensor. Safe conversion to integer:
                cls_id = int(box.cls.item() if hasattr(box.cls, "item") else box.cls[0])
                cls_name = class_names.get(cls_id, f"Class {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
    
    except Exception as e:
        import traceback
        error_msg = f"ERROR during prediction: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        return None, error_msg

# 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.Textbox(label="Detection Summary", interactive=False, lines=4)
            
    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)