services/operations_maintenance/crack_detection.py

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