Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import logging | |
| from ultralytics import YOLO | |
| # Setup logging | |
| logging.basicConfig( | |
| filename="app.log", | |
| level=logging.INFO, | |
| format="%(asctime)s - %(levelname)s - %(message)s" | |
| ) | |
| def process_frame(frame: np.ndarray) -> tuple[list[dict[str, any]], np.ndarray]: | |
| """ | |
| Process a frame using YOLO model to detect cracks and holes. | |
| Args: | |
| frame: Input frame (numpy array) | |
| Returns: | |
| tuple: List of detected items and annotated frame | |
| """ | |
| try: | |
| model = YOLO("models/yolov8m.pt") | |
| results = model(frame) | |
| detections = [] | |
| 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)} objects in frame.") | |
| return detections, frame | |
| except Exception as e: | |
| logging.error(f"Error processing frame: {str(e)}") | |
| return [], frame |