Spaces:
Runtime error
Runtime error
| 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) | |
| import random | |
| def detect_cracks(frame, model): | |
| """ | |
| Detect cracks in a frame using YOLOv8. | |
| Args: | |
| frame: Input frame (numpy array) | |
| model: YOLO model | |
| Returns: | |
| list: List of detected cracks with type, label, coordinates, confidence, and severity | |
| """ | |
| # Run YOLOv8 inference for cracks | |
| results = model(frame) | |
| detected_items = [] | |
| line_counter = 1 # Initialize counter for numbered labels | |
| 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 != "crack": # Process only cracks | |
| continue | |
| xyxy = box.xyxy[0].cpu().numpy() | |
| x_min, y_min, x_max, y_max = map(int, xyxy) | |
| # Simulate severity for cracks | |
| severity = random.choice(["low", "medium", "high"]) | |
| # Add numbered label | |
| detection_label = f"Line {line_counter} - Crack (Conf: {conf:.2f})" | |
| item = { | |
| "type": label, | |
| "label": detection_label, | |
| "confidence": conf, | |
| "coordinates": [x_min, y_min, x_max, y_max], | |
| "severity": severity | |
| } | |
| detected_items.append(item) | |
| line_counter += 1 | |
| return detected_items | |
| def detect_potholes(frame, model): | |
| """ | |
| Detect potholes in a frame using YOLOv8. | |
| Args: | |
| frame: Input frame (numpy array) | |
| model: YOLO model | |
| Returns: | |
| list: List of detected potholes with type, label, coordinates, and confidence | |
| """ | |
| # Run YOLOv8 inference for potholes | |
| results = model(frame) | |
| detected_items = [] | |
| line_counter = 1 # Initialize counter for numbered labels | |
| 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 != "pothole": # Process only potholes | |
| continue | |
| xyxy = box.xyxy[0].cpu().numpy() | |
| x_min, y_min, x_max, y_max = map(int, xyxy) | |
| # Add numbered label | |
| detection_label = f"Line {line_counter} - Pothole (Conf: {conf:.2f})" | |
| item = { | |
| "type": label, | |
| "label": detection_label, | |
| "confidence": conf, | |
| "coordinates": [x_min, y_min, x_max, y_max] | |
| } | |
| detected_items.append(item) | |
| line_counter += 1 | |
| return detected_items | |
| def detect_objects(frame, model): | |
| """ | |
| Detect objects in a frame using YOLOv8. | |
| Args: | |
| frame: Input frame (numpy array) | |
| model: YOLO model | |
| Returns: | |
| list: List of detected objects with type, label, coordinates, and confidence | |
| """ | |
| # Run YOLOv8 inference for other objects | |
| results = model(frame) | |
| detected_items = [] | |
| line_counter = 1 # Initialize counter for numbered labels | |
| 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 != "object": # Process only objects | |
| continue | |
| xyxy = box.xyxy[0].cpu().numpy() | |
| x_min, y_min, x_max, y_max = map(int, xyxy) | |
| # Add numbered label | |
| detection_label = f"Line {line_counter} - Object (Conf: {conf:.2f})" | |
| item = { | |
| "type": label, | |
| "label": detection_label, | |
| "confidence": conf, | |
| "coordinates": [x_min, y_min, x_max, y_max] | |
| } | |
| detected_items.append(item) | |
| line_counter += 1 | |
| return detected_items | |
| def detect_items_in_sequence(frame, model): | |
| """ | |
| Run crack, pothole, and object detection sequentially. | |
| Args: | |
| frame: Input frame (numpy array) | |
| model: YOLO model | |
| Returns: | |
| list: List of detected items (crack, pothole, object) | |
| """ | |
| detected_items = [] | |
| # Detect cracks first | |
| detected_items.extend(detect_cracks(frame, model)) | |
| # Detect potholes second | |
| detected_items.extend(detect_potholes(frame, model)) | |
| # Detect objects third | |
| detected_items.extend(detect_objects(frame, model)) | |
| return detected_items | |