Dina-Raslan/concrete-damage-detector

1import gradio as gr2import numpy as np3import onnxruntime as ort4from PIL import Image, ImageDraw, ImageFont5import cv26 7# ── Configuration ─────────────────────────────────────────────────────────8MODEL_PATH = "best.onnx"9INPUT_SIZE = 64010CONF_THRESHOLD = 0.411IOU_THRESHOLD = 0.4512 13CLASS_NAMES = ["crack", "spalling", "pothole"]14CLASS_COLORS = {15    "crack": (255, 0, 0),       # red16    "spalling": (255, 165, 0),  # orange17    "pothole": (255, 255, 0),   # yellow18}19 20# ── Load ONNX model once at startup ─────────────────────────────────────────21session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])22input_name = session.get_inputs()[0].name23 24 25# ── Preprocessing ────────────────────────────────────────────────────────26def preprocess(image: Image.Image):27    original_w, original_h = image.size28    resized = image.resize((INPUT_SIZE, INPUT_SIZE))29    img_array = np.array(resized).astype(np.float32) / 255.030    # HWC -> CHW31    img_array = img_array.transpose(2, 0, 1)32    # Add batch dimension33    img_array = np.expand_dims(img_array, axis=0)34    return img_array, original_w, original_h35 36 37# ── IoU + NMS ─────────────────────────────────────────────────────────────38def compute_iou(box1, box2):39    x1 = max(box1[0], box2[0])40    y1 = max(box1[1], box2[1])41    x2 = min(box1[2], box2[2])42    y2 = min(box1[3], box2[3])43    inter_area = max(0, x2 - x1) * max(0, y2 - y1)44    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])45    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])46    union_area = box1_area + box2_area - inter_area47    return inter_area / union_area if union_area > 0 else 048 49 50def non_max_suppression(detections, iou_threshold):51    detections = sorted(detections, key=lambda d: d["confidence"], reverse=True)52    keep = []53    while detections:54        best = detections.pop(0)55        keep.append(best)56        detections = [57            d for d in detections58            if compute_iou(best["bbox"], d["bbox"]) < iou_threshold59        ]60    return keep61 62 63# ── Parse YOLOv8 raw output: shape [1, 7, 8400] ────────────────────────────64# 7 = 4 (cx, cy, w, h) + 3 (class scores, no separate objectness column)65# 8400 = number of candidate detections across all scales66def parse_yolo_output(output, original_w, original_h):67    output = output[0]  # shape: [7, 8400]68    output = output.T   # transpose -> shape: [8400, 7]69 70    scale_x = original_w / INPUT_SIZE71    scale_y = original_h / INPUT_SIZE72 73    detections = []74    for row in output:75        cx, cy, w, h = row[0], row[1], row[2], row[3]76        class_scores = row[4:4 + len(CLASS_NAMES)]77        class_id = int(np.argmax(class_scores))78        confidence = float(class_scores[class_id])79 80        if confidence < CONF_THRESHOLD:81            continue82 83        cx, cy, w, h = cx * scale_x, cy * scale_y, w * scale_x, h * scale_y84        x1, y1 = cx - w / 2, cy - h / 285        x2, y2 = cx + w / 2, cy + h / 286 87        detections.append({88            "bbox": [x1, y1, x2, y2],89            "confidence": confidence,90            "class_id": class_id,91            "class_name": CLASS_NAMES[class_id],92        })93 94    return non_max_suppression(detections, IOU_THRESHOLD)95 96 97# ── Draw detections on the image ────────────────────────────────────────98def draw_detections(image: Image.Image, detections):99    image = image.copy()100    draw = ImageDraw.Draw(image)101 102    for det in detections:103        x1, y1, x2, y2 = det["bbox"]104        color = CLASS_COLORS.get(det["class_name"], (255, 255, 255))105        label = f"{det['class_name']} {det['confidence']:.2f}"106 107        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)108        text_bbox = draw.textbbox((x1, y1), label)109        draw.rectangle(text_bbox, fill=color)110        draw.text((x1, y1), label, fill=(0, 0, 0))111 112    return image113 114 115# ── Main inference function (called by Gradio) ─────────────────────────────116def detect_damage(input_image: Image.Image):117    if input_image is None:118        return None, "No image provided."119 120    input_image = input_image.convert("RGB")121    img_array, original_w, original_h = preprocess(input_image)122 123    outputs = session.run(None, {input_name: img_array})124    detections = parse_yolo_output(outputs[0], original_w, original_h)125 126    result_image = draw_detections(input_image, detections)127 128    if not detections:129        summary = "No damage detected."130    else:131        lines = [f"Found {len(detections)} detection(s):"]132        for d in detections:133            lines.append(134                f"- {d['class_name']} (confidence: {d['confidence']:.2f})"135            )136        summary = "\n".join(lines)137 138    return result_image, summary139 140 141# ── Gradio interface ────────────────────────────────────────────────────142demo = gr.Interface(143    fn=detect_damage,144    inputs=gr.Image(type="pil", label="Upload an image"),145    outputs=[146        gr.Image(type="pil", label="Detection result"),147        gr.Textbox(label="Summary"),148    ],149    title="Structural Damage Detection (Crack / Spalling / Pothole)",150    description="Upload an image of a road or concrete structure to detect cracks, spalling, and potholes using a YOLOv8 model.",151)152 153if __name__ == "__main__":154    demo.launch()