app.py

import warnings
warnings.filterwarnings('ignore')

import gradio as gr
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.gridspec import GridSpec
from ultralytics import RTDETR
from PIL import Image
import os

# ================================================
# LOAD MODELS
# ================================================
print("Loading RT-DETR...")
rtdetr_model = RTDETR('best.pt')

print("Loading MiDaS...")
midas = torch.hub.load('intel-isl/MiDaS', 'DPT_Large', trust_repo=True)
midas.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
midas.to(device)
midas_transforms = torch.hub.load('intel-isl/MiDaS', 'transforms', trust_repo=True)
transform = midas_transforms.dpt_transform
print("All models loaded!")

# ================================================
# ROAD IMAGE CHECK
# ================================================
def is_road_image(img_rgb):
    """
    Returns True only if the image looks like a road/ground surface.
    Checks for dark/grey ground tones typical of asphalt.
    """
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    h, w = gray.shape

    # Check bottom half of image (road is usually at the bottom)
    bottom = img_rgb[h//2:, :, :]
    r = bottom[:,:,0].astype(float)
    g = bottom[:,:,1].astype(float)
    b = bottom[:,:,2].astype(float)

    # Asphalt/road pixels are dark grey: low brightness, low saturation
    brightness   = np.mean(gray)
    mean_r       = np.mean(r)
    mean_g       = np.mean(g)
    mean_b       = np.mean(b)

    # Channel difference — road pixels have similar R,G,B (grey)
    channel_diff = np.mean(np.abs(r - g) + np.abs(g - b) + np.abs(r - b))

    # Road: moderate brightness, low channel variance, greyish
    is_grey_enough  = channel_diff < 60      # road is grey (not colorful)
    is_dark_enough  = brightness < 180       # road is not very bright
    is_not_sky      = mean_b < 160          # sky images have high blue
    is_not_logo     = brightness > 20        # logos/black images excluded

    return is_grey_enough and is_dark_enough and is_not_sky and is_not_logo

# ================================================
# WEATHER DETECTION
# ================================================
def detect_weather(img_rgb):
    gray          = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    brightness    = np.mean(gray)
    laplacian_var = cv2.Laplacian(gray, cv2.CV_64F).var()
    r, g, b       = img_rgb[:,:,0], img_rgb[:,:,1], img_rgb[:,:,2]
    blue_ratio    = np.mean(b) / (np.mean(r) + 1)
    dark_ratio    = np.sum(
        (r.astype(int)+g.astype(int)+b.astype(int)) < 150
    ) / (img_rgb.shape[0] * img_rgb.shape[1])
    grey_mask = (
        (np.abs(r.astype(int)-g.astype(int)) < 20) &
        (np.abs(g.astype(int)-b.astype(int)) < 20) &
        (r > 150)
    )
    fog_ratio = np.sum(grey_mask) / (img_rgb.shape[0] * img_rgb.shape[1])

    if brightness < 60:
        return 'Night',     'Low visibility - drive carefully!'
    if fog_ratio > 0.30 and laplacian_var < 200:
        return 'Foggy',     'Reduced visibility - slow down!'
    if blue_ratio > 1.10 or dark_ratio > 0.20:
        return 'Rainy/Wet', 'Wet road - risk of skidding!'
    return         'Normal', 'Good visibility conditions'

# ================================================
# SEVERITY
# ================================================
def get_severity(area_ratio, depth_score):
    score = (area_ratio * 0.5) + (depth_score * 0.5)
    if score < 0.15:
        return 'Minor',    (255, 255,   0), 'Low'
    elif score < 0.35:
        return 'Moderate', (255, 165,   0), 'Medium'
    else:
        return 'Severe',   (255,   0,   0), 'High'

def px_to_cm(px, dim): return (px / dim) * 500

# ================================================
# MAIN DETECTION FUNCTION
# ================================================
def detect_potholes(input_image):

    if input_image is None:
        return None, "Please upload an image first."

    img_rgb = np.array(input_image.convert('RGB'))
    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
    H, W    = img_rgb.shape[:2]

    # ---- ROAD CHECK ----
    if not is_road_image(img_rgb):
        fig = plt.figure(figsize=(8, 4))
        fig.patch.set_facecolor('#0d1117')
        ax = fig.add_subplot(111)
        ax.set_facecolor('#161b22')
        ax.imshow(img_rgb)
        ax.axis('off')
        ax.set_title(
            'NOT A ROAD IMAGE — Please upload a road/street photo',
            color='#ff4444', fontsize=13, fontweight='bold')
        plt.tight_layout()
        return fig, "This does not appear to be a road image.\nPlease upload a photo of a road or street for pothole detection."

    weather_label, weather_warning = detect_weather(img_rgb)

    tmp_path = '/tmp/input_image.png'
    cv2.imwrite(tmp_path, img_bgr)

    # ---- CONFIDENCE THRESHOLD = 0.5 ----
    results = rtdetr_model(tmp_path, verbose=False, conf=0.5)
    boxes   = results[0].boxes

    inp = transform(img_rgb).to(device)
    with torch.no_grad():
        dm = midas(inp)
        dm = torch.nn.functional.interpolate(
            dm.unsqueeze(1), size=(H, W),
            mode='bicubic', align_corners=False
        ).squeeze().cpu().numpy()
    depth_norm = (dm - dm.min()) / (dm.max() - dm.min() + 1e-8)

    pothole_data    = []
    severity_counts = {'Minor': 0, 'Moderate': 0, 'Severe': 0}

    if boxes is not None and len(boxes) > 0:
        for i, box in enumerate(boxes.xyxy):
            x1, y1, x2, y2 = map(int, box.tolist())
            conf = float(boxes.conf[i])

            # ---- SKIP LOW CONFIDENCE ----
            if conf < 0.5:
                continue

            bw, bh     = x2-x1, y2-y1
            area_ratio = (bw * bh) / (H * W)
            bw_cm      = px_to_cm(bw, W)
            bh_cm      = px_to_cm(bh, H)
            region     = depth_norm[y1:y2, x1:x2]
            dscore     = float(np.mean(region)) if region.size > 0 else 0.0
            depth_cm   = dscore * 15
            sev, crgb, risk = get_severity(area_ratio, dscore)
            severity_counts[sev] += 1
            cnorm = tuple(c/255 for c in crgb)
            pothole_data.append(dict(
                id=len(pothole_data)+1, severity=sev,
                confidence=conf, width_cm=bw_cm, height_cm=bh_cm,
                area_cm2=bw_cm*bh_cm, area_pct=area_ratio*100,
                depth_cm=depth_cm, risk=risk, color=cnorm,
                box=(x1,y1,x2,y2)))

    # ================================================
    # BUILD FIGURE
    # ================================================
    fig = plt.figure(figsize=(22, 14))
    fig.patch.set_facecolor('#0d1117')
    gs = GridSpec(3, 3, figure=fig, hspace=0.45, wspace=0.30)

    ax_orig   = fig.add_subplot(gs[0, 0])
    ax_depth  = fig.add_subplot(gs[0, 1:3])
    ax_detect = fig.add_subplot(gs[1, 0:3])
    ax_bar    = fig.add_subplot(gs[2, 0])
    ax_area   = fig.add_subplot(gs[2, 1])
    ax_dim    = fig.add_subplot(gs[2, 2])

    for ax in [ax_orig, ax_depth, ax_detect, ax_bar, ax_area, ax_dim]:
        ax.set_facecolor('#161b22')
        for spine in ax.spines.values():
            spine.set_edgecolor('#30363d')

    ax_orig.imshow(img_rgb)
    ax_orig.axis('off')
    ax_orig.set_title('Original Image', color='white', fontsize=11, fontweight='bold')

    ax_depth.imshow(depth_norm, cmap='plasma')
    ax_depth.axis('off')
    ax_depth.set_title('MiDaS Depth Map', color='white', fontsize=11, fontweight='bold')

    ax_detect.imshow(img_rgb)
    ax_detect.set_title(
        f'RT-DETR Detection  |  Weather: {weather_label}  |  {weather_warning}',
        color='white', fontsize=13, fontweight='bold')
    ax_detect.axis('off')

    for p in pothole_data:
        x1,y1,x2,y2 = p['box']
        bw,bh = x2-x1, y2-y1
        cnorm = p['color']
        rect = patches.FancyBboxPatch(
            (x1,y1), bw, bh, boxstyle="round,pad=2",
            linewidth=2.5, edgecolor=cnorm, facecolor=(*cnorm, 0.10))
        ax_detect.add_patch(rect)
        ax_detect.text(
            x1, max(y1-8,10),
            f'#{p["id"]} {p["severity"]} ({p["area_pct"]:.1f}%) | Conf:{p["confidence"]*100:.0f}%',
            color=cnorm, fontsize=10, fontweight='bold',
            bbox=dict(boxstyle='round,pad=2', facecolor='#0d1117',
                      alpha=0.85, edgecolor=cnorm))
        ax_depth.add_patch(patches.Rectangle(
            (x1,y1), bw, bh,
            linewidth=1.5, edgecolor='white', facecolor='none'))
        ax_depth.text(x1+2, y1+14, f'#{p["id"]}',
                      color='white', fontsize=8, fontweight='bold')

    sv = ['Minor','Moderate','Severe']
    sc = ['#ffff00','#ffa500','#ff4444']
    bars = ax_bar.bar(sv, [severity_counts[s] for s in sv],
                      color=sc, edgecolor='white', linewidth=0.5)
    ax_bar.set_title('Severity Distribution', color='white', fontsize=11, fontweight='bold')
    ax_bar.tick_params(colors='white')
    ax_bar.set_ylabel('Count', color='white')
    ax_bar.yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    for bar, val in zip(bars, [severity_counts[s] for s in sv]):
        if val > 0:
            ax_bar.text(bar.get_x()+bar.get_width()/2,
                        bar.get_height()+0.05, str(val),
                        ha='center', color='white', fontweight='bold', fontsize=12)

    if pothole_data:
        ids   = [f"#{p['id']}" for p in pothole_data]
        areas = [p['area_pct'] for p in pothole_data]
        cols  = [p['color']    for p in pothole_data]
        b2 = ax_area.bar(ids, areas, color=cols, edgecolor='white', linewidth=0.5)
        ax_area.set_title('Pothole Area (% of Image)', color='white', fontsize=11, fontweight='bold')
        ax_area.set_ylabel('Area %', color='white')
        ax_area.tick_params(colors='white')
        for bar, val in zip(b2, areas):
            ax_area.text(bar.get_x()+bar.get_width()/2,
                         bar.get_height()+0.01,
                         f'{val:.1f}%', ha='center', color='white', fontsize=9)

        x = np.arange(len(ids))
        ax_dim.bar(x-0.25, [p['width_cm']  for p in pothole_data], 0.25,
                   label='Width cm',  color='#58a6ff', edgecolor='white', linewidth=0.5)
        ax_dim.bar(x,      [p['height_cm'] for p in pothole_data], 0.25,
                   label='Height cm', color='#3fb950', edgecolor='white', linewidth=0.5)
        ax_dim.bar(x+0.25, [p['depth_cm']  for p in pothole_data], 0.25,
                   label='Depth cm',  color='#f78166', edgecolor='white', linewidth=0.5)
        ax_dim.set_xticks(x)
        ax_dim.set_xticklabels(ids, color='white', fontsize=7)
        ax_dim.tick_params(colors='white')
        ax_dim.set_title('Dimensions (cm)', color='white', fontsize=11, fontweight='bold')
        ax_dim.set_ylabel('cm', color='white')
        ax_dim.legend(facecolor='#161b22', labelcolor='white', fontsize=9)
    else:
        for ax in [ax_area, ax_dim]:
            ax.text(0.5, 0.5, 'No potholes detected',
                    ha='center', color='white',
                    transform=ax.transAxes, fontsize=11)

    fig.suptitle(
        f'POTHOLE DETECTION REPORT  |  Weather: {weather_label}  |  '
        f'{len(pothole_data)} Pothole(s) Found',
        color='white', fontsize=14, fontweight='bold',
        y=1.01, backgroundcolor='#0d1117')

    plt.tight_layout()

    report  = "POTHOLE DETECTION REPORT\n"
    report += f"{'='*50}\n"
    report += f"Weather  : {weather_label}\n"
    report += f"Warning  : {weather_warning}\n"
    report += f"Potholes : {len(pothole_data)} detected\n"
    report += f"{'='*50}\n"

    if pothole_data:
        for p in pothole_data:
            report += f"\nPothole #{p['id']}  [{p['severity'].upper()}]\n"
            report += f"  Confidence : {p['confidence']*100:.1f}%\n"
            report += f"  Width      : ~{p['width_cm']:.1f} cm\n"
            report += f"  Height     : ~{p['height_cm']:.1f} cm\n"
            report += f"  Area       : ~{p['area_cm2']:.0f} cm2 ({p['area_pct']:.1f}%)\n"
            report += f"  Est. Depth : ~{p['depth_cm']:.1f} cm\n"
            report += f"  Risk Level : {p['risk']}\n"
        report += f"\nSeverity Summary:\n"
        report += f"  Minor    : {severity_counts['Minor']}\n"
        report += f"  Moderate : {severity_counts['Moderate']}\n"
        report += f"  Severe   : {severity_counts['Severe']}\n"
    else:
        report += "\nNo potholes detected in this image!\n"

    report += f"{'='*50}"

    return fig, report

# ================================================
# GRADIO INTERFACE
# ================================================
with gr.Blocks(
    theme=gr.themes.Base(
        primary_hue="blue",
        neutral_hue="slate",
    ),
    css="""
    .gradio-container { background-color: #0d1117 !important; }
    h1 { color: #58a6ff !important; text-align: center; }
    h3 { color: #8b949e !important; text-align: center; }
    """
) as demo:

    gr.Markdown("# Pothole Detection System")
    gr.Markdown("### RT-DETR + MiDaS Depth Estimation | Severity Analysis | Weather Detection")
    gr.Markdown("---")

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Upload Road Image", height=300)
            detect_btn  = gr.Button("Detect Potholes", variant="primary", size="lg")
            gr.Markdown("""
            **Severity Legend:**
            - Minor = Small crack, Low Risk
            - Moderate = Medium hole, Medium Risk
            - Severe = Deep hole, High Risk

            **Weather Detection:**
            - Night | Foggy | Rainy/Wet | Normal

            **Model Info:**
            - RT-DETR Large
            - mAP@50: 97.28%
            - Trained on 665 pothole images
            """)

        with gr.Column(scale=2):
            output_plot   = gr.Plot(label="Detection Report")
            output_report = gr.Textbox(
                label="Detailed Text Report",
                lines=20,
                max_lines=30)

    detect_btn.click(
        fn=detect_potholes,
        inputs=input_image,
        outputs=[output_plot, output_report])

demo.launch()