| import gradio as gr |
| print("[STARTUP] Imports completed (gradio loaded)", flush=True) |
|
|
| import numpy as np |
| import onnxruntime as ort |
| from PIL import Image, ImageDraw |
|
|
| print("[STARTUP] onnxruntime, numpy, PIL loaded", flush=True) |
|
|
| |
| MODEL_PATH = "best.onnx" |
| INPUT_SIZE = 640 |
| CONF_THRESHOLD = 0.4 |
| IOU_THRESHOLD = 0.45 |
|
|
| CLASS_NAMES = [ |
| 'corrosion', 'delamination', 'cover_detachment', 'efflorescence', |
| 'crack', 'rust', 'scaling', 'spalling', 'void', |
| 'tile_crack', 'tile_delamination', 'tile_loss', 'peeling' |
| ] |
|
|
| CLASS_MAPPING = { |
| 'crack': 'Crack', |
| 'tile_crack': 'Crack', |
| 'spalling': 'Spalling', |
| 'corrosion': 'Damage', |
| 'delamination': 'Damage', |
| 'cover_detachment': 'Spalling', |
| 'efflorescence': 'Damage', |
| 'rust': 'Damage', |
| 'scaling': 'Damage', |
| 'void': 'Damage', |
| 'tile_delamination': 'Damage', |
| 'tile_loss': 'Damage', |
| 'peeling': 'Damage', |
| } |
|
|
| MERGED_COLORS = { |
| 'Crack': (255, 0, 0), |
| 'Spalling': (255, 165, 0), |
| 'Damage': (0, 0, 255), |
| } |
|
|
| print("[STARTUP] Loading ONNX model from best.onnx ...", flush=True) |
| session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"]) |
| input_name = session.get_inputs()[0].name |
| print("[STARTUP] Model loaded successfully.", flush=True) |
|
|
|
|
| |
| def preprocess(image: Image.Image): |
| original_w, original_h = image.size |
| resized = image.resize((INPUT_SIZE, INPUT_SIZE)) |
| img_array = np.array(resized).astype(np.float32) / 255.0 |
| img_array = img_array.transpose(2, 0, 1) |
| img_array = np.expand_dims(img_array, axis=0) |
| return img_array, original_w, original_h |
|
|
|
|
| |
| def compute_iou(box1, box2): |
| x1 = max(box1[0], box2[0]) |
| y1 = max(box1[1], box2[1]) |
| x2 = min(box1[2], box2[2]) |
| y2 = min(box1[3], box2[3]) |
| inter_area = max(0, x2 - x1) * max(0, y2 - y1) |
| box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) |
| box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) |
| union_area = box1_area + box2_area - inter_area |
| return inter_area / union_area if union_area > 0 else 0 |
|
|
|
|
| def non_max_suppression(detections, iou_threshold): |
| detections = sorted(detections, key=lambda d: d["confidence"], reverse=True) |
| keep = [] |
| while detections: |
| best = detections.pop(0) |
| keep.append(best) |
| detections = [d for d in detections if compute_iou(best["bbox"], d["bbox"]) < iou_threshold] |
| return keep |
|
|
|
|
| |
| def parse_yolo_output(output, original_w, original_h): |
| output = output[0] |
| output = output.T |
|
|
| scale_x = original_w / INPUT_SIZE |
| scale_y = original_h / INPUT_SIZE |
|
|
| detections = [] |
| for row in output: |
| cx, cy, w, h = row[0], row[1], row[2], row[3] |
| class_scores = row[4:4 + len(CLASS_NAMES)] |
| class_id = int(np.argmax(class_scores)) |
| confidence = float(class_scores[class_id]) |
|
|
| if confidence < CONF_THRESHOLD: |
| continue |
|
|
| cx = cx * scale_x |
| cy = cy * scale_y |
| w = w * scale_x |
| h = h * scale_y |
| x1, y1 = cx - w / 2, cy - h / 2 |
| x2, y2 = cx + w / 2, cy + h / 2 |
|
|
| original_class = CLASS_NAMES[class_id] |
| merged_class = CLASS_MAPPING[original_class] |
|
|
| detections.append({ |
| "bbox": [x1, y1, x2, y2], |
| "confidence": confidence, |
| "class_id": class_id, |
| "original_class": original_class, |
| "class_name": merged_class, |
| }) |
|
|
| return non_max_suppression(detections, IOU_THRESHOLD) |
|
|
|
|
| def draw_detections(image: Image.Image, detections): |
| image = image.copy() |
| draw = ImageDraw.Draw(image) |
| for det in detections: |
| x1, y1, x2, y2 = det["bbox"] |
| color = MERGED_COLORS.get(det["class_name"], (255, 255, 255)) |
| label = f"{det['class_name']} {det['confidence']:.2f}" |
| draw.rectangle([x1, y1, x2, y2], outline=color, width=3) |
| text_bbox = draw.textbbox((x1, y1), label) |
| draw.rectangle(text_bbox, fill=color) |
| draw.text((x1, y1), label, fill=(255, 255, 255)) |
| return image |
|
|
|
|
| def detect_damage(input_image: Image.Image): |
| print("[REQUEST] Received new image for detection", flush=True) |
|
|
| if input_image is None: |
| return None, "No image provided." |
|
|
| input_image = input_image.convert("RGB") |
| img_array, original_w, original_h = preprocess(input_image) |
|
|
| outputs = session.run(None, {input_name: img_array}) |
| detections = parse_yolo_output(outputs[0], original_w, original_h) |
|
|
| result_image = draw_detections(input_image, detections) |
|
|
| if not detections: |
| summary = "No damage detected." |
| else: |
| lines = [f"Found {len(detections)} detection(s):\n"] |
| for d in detections: |
| x1, y1, x2, y2 = d["bbox"] |
| lines.append( |
| f"- category: {d['class_name']}\n" |
| f" confidence: {d['confidence']:.2f}\n" |
| f" bbox: [{x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f}]\n" |
| ) |
| summary = "\n".join(lines) |
|
|
| print(f"[REQUEST] Done. {len(detections)} detections.", flush=True) |
| return result_image, summary |
|
|
|
|
| print("[STARTUP] Building Gradio interface ...", flush=True) |
| demo = gr.Interface( |
| fn=detect_damage, |
| inputs=gr.Image(type="pil", label="Upload an image"), |
| outputs=[ |
| gr.Image(type="pil", label="Detection result"), |
| gr.Textbox(label="Summary"), |
| ], |
| title="Urban Infrastructure Defect Detection", |
| description=( |
| "Upload an image of a building, bridge, pavement, or tiled surface to detect structural defects. " |
| "Detections are grouped into three categories: Crack (red), Spalling (orange), and Damage (blue). " |
| ), |
| ) |
|
|
| print("[STARTUP] Calling demo.launch() now ....", flush=True) |
|
|
| if __name__ == "__main__": |
| demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False) |
