Spaces:
Runtime error
Runtime error
File size: 1,900 Bytes
91cad63 c45c8a6 15507ed 91cad63 15507ed c45c8a6 15507ed c45c8a6 15507ed c45c8a6 15507ed |
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 |
import cv2
import numpy as np
import logging
from typing import List, Dict, Any, Tuple
from ultralytics import YOLO
# Setup logging
logging.basicConfig(
filename="app.log",
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s"
)
def detect_cracks_and_holes(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
"""
Detect cracks and holes in a frame using YOLO.
Args:
frame: Input frame (numpy array)
Returns:
tuple: List of detected items (dicts with type, box, severity, etc.) and the input frame
"""
try:
# Load YOLO model
model = YOLO("models/yolov8m.pt")
# Run inference
results = model(frame)
detections: List[Dict[str, Any]] = []
line_counter = 1
for result in results:
boxes = result.boxes.xyxy.cpu().numpy()
confidences = result.boxes.conf.cpu().numpy()
classes = result.boxes.cls.cpu().numpy()
for box, conf, cls in zip(boxes, confidences, classes):
x_min, y_min, x_max, y_max = map(int, box)
type_ = "crack" if cls == 0 else "hole" # Assume class 0=crack, 1=hole
severity = "Severe" if conf > 0.8 else "Moderate" if conf > 0.5 else "Minor"
label = f"Line {line_counter} - {type_.capitalize()} ({severity})"
detections.append({
"type": type_,
"label": label,
"box": [x_min, y_min, x_max, y_max],
"severity": severity,
"confidence": float(conf)
})
line_counter += 1
logging.info(f"Detected {len(detections)} cracks/holes in frame.")
return detections, frame
except Exception as e:
logging.error(f"Error detecting cracks/holes: {str(e)}")
return [], frame |