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

# Load your optimized model
# Make sure 'best_optimized_model.pt' is in the same directory on Hugging Face
model = YOLO("best_optimized_model.pt")

def detect_potholes(image, conf_threshold):
    if image is None:
        return None, "No image uploaded."
        
    # Run inference behind the scenes
    results = model(image, conf=conf_threshold)
    
    # Count the hidden detections
    boxes = results[0].boxes
    box_count = len(boxes)
    
    # Rule: If less than 2 boxes are detected, treat as an anomaly/clean road
    if box_count < 2:
        status_text = "✅ Road looks safe! No potholes detected."
        # Returns the original image completely clean of any bounding boxes
        return image, status_text
        
    # Rule: If 2 or more boxes are found, confirm pothole presence
    status_text = "⚠️ Warning: Potholes Detected!"
    
    # Returns the original image untouched, keeping the layout clean as requested
    return image, status_text

# Define the user interface layout
interface = gr.Interface(
    fn=detect_potholes,
    inputs=[
        gr.Image(type="pil", label="Upload Road Image"),
        gr.Slider(
            minimum=0.01, 
            maximum=1.0, 
            value=0.25, 
            step=0.01, 
            label="Confidence Threshold"
        )
    ],
    outputs=[
        gr.Image(type="pil", label="Road Image Preview"),
        gr.Textbox(label="System Status Analysis")
    ],
    title="Heuristic Pothole Detector (Clean Output Mode)",
    description=(
        "This version evaluates pothole presence using a custom multi-point validation heuristic. "
        "The output image remains clean without any bounding box overlays, and system alerts hide "
        "the raw numerical box counts."
    )
)

if __name__ == "__main__":
    interface.launch()