Update app.py

app.py

@@ -1,144 +1,40 @@
 import gradio as gr
-import numpy as np
-import onnxruntime as ort
-from PIL import Image, ImageDraw, ImageFont
-import cv2
+from ultralytics import YOLO
+from PIL import Image
 
-# ── Configuration ─────────────────────────────────────────────────────────
-MODEL_PATH = "best.onnx"
-INPUT_SIZE = 640
-CONF_THRESHOLD = 0.4
-IOU_THRESHOLD = 0.45
+# ── Load the model once at startup ──────────────────────────────────────
+model = YOLO("best.pt")
 
-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)
+    results = model.predict(
+        source=input_image,
+        imgsz=640,
+        conf=0.4,
+        iou=0.45,
+    )
 
-    result_image = draw_detections(input_image, detections)
+    result = results[0]
+    result_image = Image.fromarray(result.plot()[:, :, ::-1])  # BGR -> RGB
 
-    if not detections:
+    detections = result.boxes
+    if detections is None or len(detections) == 0:
         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})"
-            )
+        for box in detections:
+            class_id = int(box.cls[0])
+            confidence = float(box.conf[0])
+            class_name = model.names[class_id]
+            lines.append(f"- {class_name} (confidence: {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"),
@@ -151,4 +47,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(share=True)