Spaces:

Dina-Raslan
/

concrete-damage-detector

Paused

Create app.py

58aa9b1 verified 4 months ago

5.94 kB

import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont
import cv2

# ── Configuration ─────────────────────────────────────────────────────────
MODEL_PATH = "best.onnx"
INPUT_SIZE = 640
CONF_THRESHOLD = 0.4
IOU_THRESHOLD = 0.45

CLASS_NAMES = ["crack", "spalling", "pothole"]
CLASS_COLORS = {
    "crack": (255, 0, 0),       # red
    "spalling": (255, 165, 0),  # orange
    "pothole": (255, 255, 0),   # yellow
}

# ── Load ONNX model once at startup ─────────────────────────────────────────
session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name


# ── Preprocessing ────────────────────────────────────────────────────────
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
    # HWC -> CHW
    img_array = img_array.transpose(2, 0, 1)
    # Add batch dimension
    img_array = np.expand_dims(img_array, axis=0)
    return img_array, original_w, original_h


# ── IoU + NMS ─────────────────────────────────────────────────────────────
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


# ── Parse YOLOv8 raw output: shape [1, 7, 8400] ────────────────────────────
# 7 = 4 (cx, cy, w, h) + 3 (class scores, no separate objectness column)
# 8400 = number of candidate detections across all scales
def parse_yolo_output(output, original_w, original_h):
    output = output[0]  # shape: [7, 8400]
    output = output.T   # transpose -> shape: [8400, 7]

    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, cy, w, h = cx * scale_x, cy * scale_y, w * scale_x, h * scale_y
        x1, y1 = cx - w / 2, cy - h / 2
        x2, y2 = cx + w / 2, cy + h / 2

        detections.append({
            "bbox": [x1, y1, x2, y2],
            "confidence": confidence,
            "class_id": class_id,
            "class_name": CLASS_NAMES[class_id],
        })

    return non_max_suppression(detections, IOU_THRESHOLD)


# ── Draw detections on the image ────────────────────────────────────────
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 = CLASS_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=(0, 0, 0))

    return image


# ── Main inference function (called by Gradio) ─────────────────────────────
def detect_damage(input_image: Image.Image):
    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):"]
        for d in detections:
            lines.append(
                f"- {d['class_name']} (confidence: {d['confidence']:.2f})"
            )
        summary = "\n".join(lines)

    return result_image, summary


# ── Gradio interface ────────────────────────────────────────────────────
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="Structural Damage Detection (Crack / Spalling / Pothole)",
    description="Upload an image of a road or concrete structure to detect cracks, spalling, and potholes using a YOLOv8 model.",
)

if __name__ == "__main__":
    demo.launch()

	import gradio as gr
	import numpy as np
	import onnxruntime as ort
	from PIL import Image, ImageDraw, ImageFont
	import cv2

	# ── Configuration ─────────────────────────────────────────────────────────
	MODEL_PATH = "best.onnx"
	INPUT_SIZE = 640
	CONF_THRESHOLD = 0.4
	IOU_THRESHOLD = 0.45

	CLASS_NAMES = ["crack", "spalling", "pothole"]
	CLASS_COLORS = {
	"crack": (255, 0, 0), # red
	"spalling": (255, 165, 0), # orange
	"pothole": (255, 255, 0), # yellow
	}

	# ── Load ONNX model once at startup ─────────────────────────────────────────
	session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
	input_name = session.get_inputs()[0].name


	# ── Preprocessing ────────────────────────────────────────────────────────
	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
	# HWC -> CHW
	img_array = img_array.transpose(2, 0, 1)
	# Add batch dimension
	img_array = np.expand_dims(img_array, axis=0)
	return img_array, original_w, original_h


	# ── IoU + NMS ─────────────────────────────────────────────────────────────
	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


	# ── Parse YOLOv8 raw output: shape [1, 7, 8400] ────────────────────────────
	# 7 = 4 (cx, cy, w, h) + 3 (class scores, no separate objectness column)
	# 8400 = number of candidate detections across all scales
	def parse_yolo_output(output, original_w, original_h):
	output = output[0] # shape: [7, 8400]
	output = output.T # transpose -> shape: [8400, 7]

	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, cy, w, h = cx * scale_x, cy * scale_y, w * scale_x, h * scale_y
	x1, y1 = cx - w / 2, cy - h / 2
	x2, y2 = cx + w / 2, cy + h / 2

	detections.append({
	"bbox": [x1, y1, x2, y2],
	"confidence": confidence,
	"class_id": class_id,
	"class_name": CLASS_NAMES[class_id],
	})

	return non_max_suppression(detections, IOU_THRESHOLD)


	# ── Draw detections on the image ────────────────────────────────────────
	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 = CLASS_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=(0, 0, 0))

	return image


	# ── Main inference function (called by Gradio) ─────────────────────────────
	def detect_damage(input_image: Image.Image):
	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):"]
	for d in detections:
	lines.append(
	f"- {d['class_name']} (confidence: {d['confidence']:.2f})"
	)
	summary = "\n".join(lines)

	return result_image, summary


	# ── Gradio interface ────────────────────────────────────────────────────
	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="Structural Damage Detection (Crack / Spalling / Pothole)",
	description="Upload an image of a road or concrete structure to detect cracks, spalling, and potholes using a YOLOv8 model.",
	)

	if __name__ == "__main__":
	demo.launch()