Pothole YOLO pipeline: single-image mode, modular detector, red box visualization

app/main.py

@@ -222,8 +222,8 @@ def me(user: Optional[User] = Depends(get_current_user)):
 @app.post("/api/detect")
 async def detect(
     request: Request,
-    before: UploadFile = File(...),
-    after: UploadFile = File(...),
+    before: Optional[UploadFile] = File(None),
+    after: Optional[UploadFile] = File(None),
     method: str = Form("AI-Based Deep Learning"),
     detection_type: str = Form("change_detection"),
     landslide_model: str = Form("Rule-Based v1"),
@@ -252,23 +252,42 @@ async def detect(
     if not user:
         raise HTTPException(status_code=401, detail="Login required")
     MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20 MB
-    try:
-        before_bytes = await before.read()
-        after_bytes = await after.read()
-        if len(before_bytes) > MAX_UPLOAD_BYTES or len(after_bytes) > MAX_UPLOAD_BYTES:
-            raise HTTPException(status_code=400, detail="Image too large (max 20 MB)")
-        before_pil = Image.open(io.BytesIO(before_bytes)).convert("RGB")
-        after_pil = Image.open(io.BytesIO(after_bytes)).convert("RGB")
-        del before_bytes, after_bytes
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Invalid image: {e}")
+    detection_type = (detection_type or "change_detection").strip().lower()
+
+    def _read_upload(upload: Optional[UploadFile], field_name: str):
+        if upload is None:
+            raise HTTPException(status_code=400, detail=f"{field_name} image is required")
+        raw = None
+        try:
+            raw = upload.file.read()
+            if raw is None or len(raw) == 0:
+                raise HTTPException(status_code=400, detail=f"{field_name} image is empty")
+            if len(raw) > MAX_UPLOAD_BYTES:
+                raise HTTPException(status_code=400, detail="Image too large (max 20 MB)")
+            return Image.open(io.BytesIO(raw)).convert("RGB")
+        except HTTPException:
+            raise
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid {field_name} image: {e}")
+        finally:
+            try:
+                if raw is not None:
+                    del raw
+            except Exception:
+                pass
+
+    if detection_type == "pothole_detection":
+        # Single-image mode: use after if present, else before.
+        primary = after if after is not None else before
+        after_pil = _read_upload(primary, "road")
+        before_pil = after_pil
+    else:
+        before_pil = _read_upload(before, "before")
+        after_pil = _read_upload(after, "after")
     detection_sensitivity = max(0.0, min(1.0, float(detection_sensitivity)))
     if min_region_area is not None:
         min_region_area = int(max(50, min(10000, min_region_area)))
 
-    detection_type = (detection_type or "change_detection").strip().lower()
     if detection_type == "landslide_detection":
         from .landslide_engine import run_landslide_detection
         method = f"Landslide - {landslide_model}"

app/pothole_detection/__init__.py

@@ -0,0 +1,2 @@
+from .pothole_detector import PotholeDetector
+

app/pothole_detection/inference.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import List, Dict
+
+import numpy as np
+
+
+def run_pothole_inference(
+    model,
+    image_bgr: np.ndarray,
+    conf_threshold: float = 0.25,
+    iou_threshold: float = 0.45,
+) -> List[Dict]:
+    """
+    Run YOLO inference and normalize predictions to a simple list format.
+    """
+    results = model.predict(
+        source=image_bgr,
+        conf=conf_threshold,
+        iou=iou_threshold,
+        verbose=False,
+    )
+    preds: List[Dict] = []
+    if not results:
+        return preds
+
+    r = results[0]
+    names = getattr(r, "names", {}) or {}
+    boxes = getattr(r, "boxes", None)
+    if boxes is None:
+        return preds
+
+    xyxy = boxes.xyxy.cpu().numpy() if hasattr(boxes.xyxy, "cpu") else boxes.xyxy
+    confs = boxes.conf.cpu().numpy() if hasattr(boxes.conf, "cpu") else boxes.conf
+    clss = boxes.cls.cpu().numpy() if hasattr(boxes.cls, "cpu") else boxes.cls
+
+    for i in range(len(xyxy)):
+        x1, y1, x2, y2 = [int(v) for v in xyxy[i]]
+        confidence = float(confs[i])
+        cls_id = int(clss[i]) if clss is not None else 0
+        cls_name = names.get(cls_id, "pothole")
+        preds.append(
+            {
+                "bbox": [x1, y1, x2, y2],
+                "confidence": confidence,
+                "class_id": cls_id,
+                "class_name": str(cls_name),
+            }
+        )
+
+    return preds
+

app/pothole_detection/model_loader.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+
+
+@lru_cache(maxsize=1)
+def get_yolo_model():
+    """
+    Lazy-load Ultralytics YOLO model once per process.
+
+    Env:
+    - POTHOLE_MODEL_PATH: local path or model name (default: yolov8n.pt)
+    """
+    model_path = os.environ.get("POTHOLE_MODEL_PATH", "yolov8n.pt").strip() or "yolov8n.pt"
+    from ultralytics import YOLO
+    return YOLO(model_path)
+

app/pothole_detection/pothole_detector.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import Dict, Any, List
+
+import cv2
+import numpy as np
+
+from .model_loader import get_yolo_model
+from .inference import run_pothole_inference
+from .visualization import draw_pothole_boxes
+
+
+class PotholeDetector:
+    """
+    Modular pothole detector:
+    - preprocessing
+    - model inference
+    - post-processing
+    - visualization
+    """
+
+    def __init__(self, conf_threshold: float = 0.25, iou_threshold: float = 0.45):
+        self.conf_threshold = float(conf_threshold)
+        self.iou_threshold = float(iou_threshold)
+        self.model = get_yolo_model()
+
+    def preprocess(self, image_bgr: np.ndarray) -> np.ndarray:
+        # Lightweight denoise for road textures
+        return cv2.bilateralFilter(image_bgr, 5, 35, 35)
+
+    def infer(self, image_bgr: np.ndarray) -> List[Dict[str, Any]]:
+        return run_pothole_inference(
+            self.model,
+            image_bgr,
+            conf_threshold=self.conf_threshold,
+            iou_threshold=self.iou_threshold,
+        )
+
+    def postprocess(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        # Keep all detections; custom filtering can be added here.
+        return detections
+
+    def visualize(self, image_bgr: np.ndarray, detections: List[Dict[str, Any]]) -> np.ndarray:
+        return draw_pothole_boxes(image_bgr, detections)
+
+    def run(self, image_bgr: np.ndarray):
+        prep = self.preprocess(image_bgr)
+        detections = self.infer(prep)
+        detections = self.postprocess(detections)
+        vis = self.visualize(image_bgr, detections)
+        return detections, vis
+

app/pothole_detection/visualization.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import List, Dict
+
+import cv2
+import numpy as np
+
+
+def draw_pothole_boxes(image_bgr: np.ndarray, detections: List[Dict]) -> np.ndarray:
+    """
+    Draw red bounding boxes with confidence labels.
+    """
+    out = image_bgr.copy()
+    for det in detections:
+        x1, y1, x2, y2 = det["bbox"]
+        conf = float(det.get("confidence", 0.0))
+        label = f"pothole {conf:.2f}"
+
+        # Red box (BGR)
+        cv2.rectangle(out, (x1, y1), (x2, y2), (0, 0, 255), 2)
+
+        # Label background
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
+        y_text = max(16, y1 - 6)
+        cv2.rectangle(out, (x1, y_text - th - 6), (x1 + tw + 8, y_text + 2), (0, 0, 255), -1)
+        cv2.putText(out, label, (x1 + 4, y_text - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
+    return out
+

app/pothole_engine.py

@@ -1,18 +1,19 @@
 """
-Pothole / road damage detection starter engine.
+Pothole / road damage detection engine (YOLO-ready).
 
-Goal: separate pipeline that can evolve to a real model (YOLO/Mask R-CNN/SegFormer).
-This initial version is a CPU-friendly heuristic detector designed for vehicle/drone imagery.
-
-Notes:
-- Satellite imagery is generally too coarse for potholes unless very high resolution (<10 cm/px).
-- Vehicle camera or low-altitude drone is the realistic input for pothole detection.
+Uses modular pipeline under app/pothole_detection:
+- model_loader.py
+- inference.py
+- visualization.py
+- pothole_detector.py
 """
 from __future__ import annotations
 
-import cv2
 import numpy as np
 from PIL import Image
+import cv2
+
+from .pothole_detection import PotholeDetector
 
 
 def _preprocess(image: Image.Image, max_size: int = 1600) -> np.ndarray:
@@ -33,156 +34,83 @@ def _norm01(x: np.ndarray) -> np.ndarray:
     return (x - lo) / (hi - lo)
 
 
-def _road_texture_response(gray: np.ndarray) -> np.ndarray:
-    # Potholes often appear as dark regions with sharp boundaries + rough texture.
-    blur = cv2.GaussianBlur(gray, (5, 5), 0)
-    lap = cv2.Laplacian(blur, cv2.CV_32F, ksize=3)
-    rough = cv2.GaussianBlur(np.abs(lap), (7, 7), 0)
-    return _norm01(rough)
-
-
-def _shadow_score(gray: np.ndarray) -> np.ndarray:
-    # Darker-than-local background regions.
-    local = cv2.GaussianBlur(gray, (31, 31), 0)
-    diff = np.clip((local - gray).astype(np.float32), 0, None)
-    return _norm01(diff)
-
-
-def _edge_score(gray: np.ndarray) -> np.ndarray:
-    med = float(np.median(gray))
-    t1 = int(max(0, 0.66 * med))
-    t2 = int(min(255, 1.33 * med))
-    edges = cv2.Canny(gray, t1, t2)
-    edges = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=1)
-    return edges.astype(np.float32) / 255.0
-
-
-def _clean(mask: np.ndarray) -> np.ndarray:
-    m = mask.copy()
-    m = cv2.medianBlur(m, 5)
-    k_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-    k_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
-    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, k_open)
-    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k_close)
-    return m
+def run_pothole_detection(
+    before_pil: Image.Image,
+    after_pil: Image.Image,
+    model_name: str = "Rule-Based v1",
+    detection_sensitivity: float = 0.6,
+    min_region_area: int | None = None,
+):
+    """
+    Current UI uses (before, after) upload. For potholes, we treat the provided road
+    image as the target and run YOLO-style detection.
+    """
+    img = _preprocess(after_pil)
+    # Ultralytics model expects BGR ndarray from OpenCV style pipeline.
+    bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 
+    # Sensitivity maps to confidence threshold inversely.
+    sens = float(np.clip(detection_sensitivity, 0.0, 1.0))
+    conf_thr = float(np.clip(0.45 - (sens - 0.5) * 0.35, 0.10, 0.70))
+    iou_thr = 0.45
+    detector = PotholeDetector(conf_threshold=conf_thr, iou_threshold=iou_thr)
+    detections, vis_bgr = detector.run(bgr)
+    result = cv2.cvtColor(vis_bgr, cv2.COLOR_BGR2RGB)
 
-def _extract_regions(mask: np.ndarray, min_area: int = 220):
-    n, labels, stats, cents = cv2.connectedComponentsWithStats(mask, connectivity=8)
-    h, w = mask.shape[:2]
-    img_area = h * w
-    regs = []
+    regions = []
     rid = 0
-    for i in range(1, n):
-        area = int(stats[i, cv2.CC_STAT_AREA])
-        if area < min_area:
+    for d in detections:
+        x1, y1, x2, y2 = d["bbox"]
+        w = max(1, x2 - x1)
+        h = max(1, y2 - y1)
+        area = int(w * h)
+        if min_region_area is not None and area < int(min_region_area):
             continue
-        x = int(stats[i, cv2.CC_STAT_LEFT])
-        y = int(stats[i, cv2.CC_STAT_TOP])
-        bw = int(stats[i, cv2.CC_STAT_WIDTH])
-        bh = int(stats[i, cv2.CC_STAT_HEIGHT])
-        if bw * bh > img_area * 0.25:
-            continue
-        ar = max(bw, bh) / max(1, min(bw, bh))
-        if ar > 6.0:
-            continue
-        cx, cy = cents[i]
-        fill = area / max(1, bw * bh)
-        conf = float(np.clip(0.25 + fill * 0.7, 0.25, 0.95))
-        sev = "minor"
-        if area / img_area > 0.01:
-            sev = "major"
-        elif area / img_area > 0.003:
-            sev = "moderate"
         rid += 1
-        regs.append(
+        conf = float(d.get("confidence", 0.0))
+        severity = "minor"
+        area_ratio = area / max(1, img.shape[0] * img.shape[1])
+        if area_ratio > 0.01:
+            severity = "major"
+        elif area_ratio > 0.003:
+            severity = "moderate"
+        regions.append(
             {
                 "id": rid,
                 "area": area,
-                "bbox": (x, y, bw, bh),
-                "center": (int(cx), int(cy)),
+                "bbox": (int(x1), int(y1), int(w), int(h)),
+                "center": (int(x1 + w // 2), int(y1 + h // 2)),
                 "object_type": "Pothole / Road Damage",
                 "confidence": conf,
-                "severity": sev,
-                "sub_type": "Pothole",
+                "severity": severity,
+                "sub_type": str(d.get("class_name", "pothole")),
                 "sub_type_confidence": conf,
                 "estimated_stories": None,
                 "estimated_height_m": None,
                 "construction_stage": None,
             }
         )
-    regs.sort(key=lambda r: r["area"], reverse=True)
-    return regs[:80]
-
-
-def _visualize(img: np.ndarray, mask: np.ndarray, regions: list[dict]) -> np.ndarray:
-    out = img.copy().astype(np.float32)
-    m = (mask > 127).astype(np.float32)
-    # Orange overlay for road damage
-    layer = np.zeros_like(out)
-    layer[:, :, 0] = 255
-    layer[:, :, 1] = 165
-    alpha = 0.35
-    for c in range(3):
-        out[:, :, c] = out[:, :, c] * (1 - m * alpha) + layer[:, :, c] * (m * alpha)
-    vis = np.clip(out, 0, 255).astype(np.uint8)
-    for r in regions:
-        x, y, w, h = r["bbox"]
-        cv2.rectangle(vis, (x, y), (x + w, y + h), (0, 140, 255), 2)
-        cv2.putText(vis, str(r["id"]), (x + 4, max(14, y - 4)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
-    return vis
-
-
-def run_pothole_detection(
-    before_pil: Image.Image,
-    after_pil: Image.Image,
-    model_name: str = "Rule-Based v1",
-    detection_sensitivity: float = 0.6,
-    min_region_area: int | None = None,
-):
-    """
-    Current UI uses (before, after) upload. For potholes, we treat the *after* image as
-    the road image and ignore the before image.
-    """
-    img = _preprocess(after_pil)
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-
-    rough = _road_texture_response(gray)
-    shadow = _shadow_score(gray)
-    edges = _edge_score(gray)
-
-    fused = 0.45 * shadow + 0.35 * rough + 0.20 * edges
-    fused = cv2.GaussianBlur(fused.astype(np.float32), (7, 7), 0)
-
-    sens = float(np.clip(detection_sensitivity, 0.0, 1.0))
-    q = float(np.clip(0.975 - (sens - 0.5) * 0.10, 0.85, 0.985))
-    thr = float(np.quantile(fused, q))
-    mask = (fused >= thr).astype(np.uint8) * 255
-    mask = _clean(mask)
-
-    if min_region_area is None:
-        min_region_area = int(max(150, min(1200, mask.shape[0] * mask.shape[1] * 0.00005)))
-    regions = _extract_regions(mask, min_area=int(min_region_area))
-    result = _visualize(img, mask, regions)
 
-    total = int(mask.shape[0] * mask.shape[1])
-    changed = int(np.sum(mask > 127))
+    total = int(img.shape[0] * img.shape[1])
+    changed = int(sum(r["area"] for r in regions))
     stats = {
         "total_pixels": total,
         "changed_pixels": changed,
         "unchanged_pixels": total - changed,
         "change_percentage": (changed / total * 100.0) if total else 0.0,
-        "image_width": mask.shape[1],
-        "image_height": mask.shape[0],
+        "image_width": img.shape[1],
+        "image_height": img.shape[0],
         "threshold_debug": {
             "method": f"Pothole Detection ({model_name})",
-            "threshold_used": int(np.clip(thr * 255.0, 0, 255)),
-            "threshold_percentile_q": q,
+            "threshold_used": None,
+            "confidence_threshold": conf_thr,
+            "iou_threshold": iou_thr,
             "sensitivity": sens,
+            "detected_boxes": len(regions),
         },
         "params": {
             "detection_sensitivity": sens,
-            "min_region_area": int(min_region_area),
+            "min_region_area": int(min_region_area) if min_region_area is not None else None,
             "model_name": model_name,
             "input": "after_only",
         },

requirements.txt

@@ -10,3 +10,4 @@ numpy>=1.24.0
 opencv-python-headless>=4.8.0
 scikit-learn>=1.3.0
 requests>=2.28.0
+ultralytics>=8.2.0

static/js/app.js

@@ -251,12 +251,23 @@ setupUploadZone('file-after', 'name-after', 'zone-after', 'preview-after');
   function refresh() {
     const isLandslide = typeSel.value === 'landslide_detection';
     const isPothole = typeSel.value === 'pothole_detection';
+    const beforeZone = document.getElementById('zone-before');
+    const beforeInput = document.getElementById('file-before');
+    const beforeName = document.getElementById('name-before');
     if (landslideGroup) landslideGroup.classList.toggle('hidden', !isLandslide);
     if (potholeGroup) potholeGroup.classList.toggle('hidden', !isPothole);
     const hideCore = isLandslide || isPothole;
     if (methodGroup) methodGroup.classList.toggle('hidden', hideCore);
     if (regGroup) regGroup.classList.toggle('hidden', hideCore);
     if (normGroup) normGroup.classList.toggle('hidden', hideCore);
+    // Pothole mode uses a single image upload (after image).
+    if (beforeZone) beforeZone.classList.toggle('hidden', isPothole);
+    if (isPothole && beforeInput) {
+      beforeInput.value = '';
+      if (beforeName) beforeName.textContent = 'No file chosen';
+      const prev = document.getElementById('preview-before');
+      if (prev) prev.classList.add('hidden');
+    }
   }
 
   typeSel.addEventListener('change', refresh);
@@ -470,9 +481,15 @@ function stopDetectionProgress(success) {
 document.getElementById('form-detect')?.addEventListener('submit', async (e) => {
   e.preventDefault();
   hideError('dashboard-error');
+  const detectionType = document.getElementById('detect-type')?.value || 'change_detection';
   const before = document.getElementById('file-before').files?.[0];
   const after = document.getElementById('file-after').files?.[0];
-  if (!before || !after) {
+  if (detectionType === 'pothole_detection') {
+    if (!after && !before) {
+      showError('dashboard-error', 'Please upload one road image for pothole detection.');
+      return;
+    }
+  } else if (!before || !after) {
     showError('dashboard-error', 'Please select both before and after images.');
     return;
   }
@@ -485,9 +502,8 @@ document.getElementById('form-detect')?.addEventListener('submit', async (e) =>
 
   const token = getToken();
   const form = new FormData();
-  form.append('before', before);
-  form.append('after', after);
-  const detectionType = document.getElementById('detect-type')?.value || 'change_detection';
+  if (before) form.append('before', before);
+  if (after) form.append('after', after);
   form.append('detection_type', detectionType);
   form.append('method', document.getElementById('detect-method').value);
   if (detectionType === 'landslide_detection') {

templates/index.html

@@ -401,6 +401,6 @@
     </div>
   </div>
 
-  <script src="/static/js/app.js?v=31"></script>
+  <script src="/static/js/app.js?v=32"></script>
 </body>
 </html>