Spaces:
Sleeping
Sleeping
| 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() |