mrshibly's picture
Upload app.py
7f47c69 verified
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
# Load the trained model
try:
model = YOLO('best.pt')
except:
print("Warning: 'best.pt' not found. Downloading yolov8n.pt for demonstration.")
model = YOLO('yolov8n.pt')
def detect_potholes(image):
"""
Function to perform inference on an image.
"""
if image is None:
return None
# Run inference
# verbose=False reduces log clutter
results = model(image, verbose=False)
# Plot results
# results[0].plot() returns a BGR numpy array
res_plotted = results[0].plot()
# Convert BGR to RGB for Gradio
res_rgb = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
return res_rgb
# CSS to ensure the video isn't mirrored (good for back cameras)
css = """
video { transform: scaleX(1) !important; }
"""
# Create the Gradio Interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🕳️ Real-Time Pothole Detection System")
gr.Markdown("Deploy this on Hugging Face Spaces. Switch tabs for different modes.")
with gr.Tabs():
# TAB 1: Real-Time Stream
with gr.Tab("📹 Live Pothole Detection"):
gr.Markdown("**Use this tab for continuous detection (Video Stream)**")
with gr.Row():
with gr.Column():
# 'sources=["webcam"]' and NO 'mirror_webcam' argument (handled by CSS)
stream_input = gr.Image(sources=["webcam"], label="Live Camera Feed", interactive=True)
with gr.Column():
stream_output = gr.Image(label="Live Detection Output")
# Continuous stream event
stream_input.stream(fn=detect_potholes, inputs=stream_input, outputs=stream_output, show_progress=False)
# TAB 2: Upload or Capture
with gr.Tab("📷 Upload / Take Photo"):
gr.Markdown("**Use this tab to upload an image or take a single snapshot.**")
with gr.Row():
with gr.Column():
# Sources allow both upload and webcam snapshot
static_input = gr.Image(sources=["upload", "webcam"], label="Upload or Snap Photo", type="numpy")
detect_btn = gr.Button("Detect Potholes", variant="primary")
with gr.Column():
static_output = gr.Image(label="Processed Image")
# Button click event
detect_btn.click(fn=detect_potholes, inputs=static_input, outputs=static_output)
# Automatic detection on upload change (optional, but good UX)
static_input.change(fn=detect_potholes, inputs=static_input, outputs=static_output)
# Launch
if __name__ == "__main__":
demo.launch()