Spaces:
Runtime error
Runtime error
File size: 4,745 Bytes
abce474 13105d4 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 fd3b717 13105d4 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 fd3b717 4991932 |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
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
|