Spaces:
Runtime error
Runtime error
File size: 1,832 Bytes
abce474 13105d4 fd3b717 13105d4 fd3b717 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import cv2
import numpy as np
from ultralytics import YOLO
import os
import random
# Load YOLOv8m-seg model for crack detection
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "../models/yolov8m-seg.pt")
model = YOLO(MODEL_PATH)
def detect_cracks_and_objects(frame):
"""
Detect cracks and other objects in a frame using YOLOv8m-seg.
Args:
frame: Input frame (numpy array)
Returns:
list: List of detected items with type, label, coordinates, confidence, and severity
"""
# Run YOLOv8 inference
results = model(frame)
detected_items = []
line_counter = 1 # Initialize counter for numbered labels
# Process detections
for r in results:
for box in r.boxes:
conf = float(box.conf[0])
if conf < 0.5:
continue
cls = int(box.cls[0])
label = model.names[cls]
if label not in ["crack", "pothole", "object"]: # Assuming these classes exist
continue
xyxy = box.xyxy[0].cpu().numpy()
x_min, y_min, x_max, y_max = map(int, xyxy)
# Simulate severity for cracks
severity = None
if label == "crack":
severity = random.choice(["low", "medium", "high"])
# Add numbered label
detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
item = {
"type": label,
"label": detection_label,
"confidence": conf,
"coordinates": [x_min, y_min, x_max, y_max]
}
if severity:
item["severity"] = severity
detected_items.append(item)
line_counter += 1
return detected_items |