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