lokesh341 commited on
Commit
91cad63
·
1 Parent(s): 549a268

Create crack_detection.py

Browse files
services/operations_maintenance/crack_detection.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # services/operations_maintenance/crack_detection.py
2
+ import cv2
3
+ import numpy as np
4
+ from ultralytics import YOLO
5
+ import os
6
+
7
+ # Load YOLOv8m-seg model
8
+ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
9
+ MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m-seg.pt")
10
+ model = YOLO(MODEL_PATH)
11
+
12
+ def detect_cracks(frame):
13
+ """
14
+ Detect cracks in a frame using YOLOv8m-seg.
15
+ Args:
16
+ frame: Input frame (numpy array)
17
+ Returns:
18
+ dict: Detection results with numbered labels
19
+ numpy array: Annotated frame
20
+ """
21
+ # Run YOLOv8 inference
22
+ results = model(frame)
23
+
24
+ detections = []
25
+ line_counter = 1 # Initialize counter for numbered labels
26
+
27
+ # Process detections
28
+ for r in results:
29
+ for box in r.boxes:
30
+ conf = float(box.conf[0])
31
+ if conf < 0.5:
32
+ continue
33
+ cls = int(box.cls[0])
34
+ label = model.names[cls]
35
+ if label != "crack": # Assuming "crack" class exists
36
+ continue
37
+ xyxy = box.xyxy[0].cpu().numpy()
38
+ x_min, y_min, x_max, y_max = map(int, xyxy)
39
+
40
+ # Add numbered label
41
+ detection_label = f"Line {line_counter} - Crack (Conf: {conf:.2f})"
42
+ detections.append({
43
+ "label": detection_label,
44
+ "confidence": conf,
45
+ "coordinates": [x_min, y_min, x_max, y_max]
46
+ })
47
+
48
+ # Draw bounding box and label
49
+ color = (255, 255, 0) # Yellow for cracks
50
+ cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
51
+ cv2.putText(frame, detection_label, (x_min, y_min - 10),
52
+ cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
53
+
54
+ line_counter += 1
55
+
56
+ return {"detections": detections, "frame": frame}
57
+
58
+ def process_cracks(frame):
59
+ """
60
+ Wrapper function for integration with app.py.
61
+ """
62
+ result = detect_cracks(frame)
63
+ return result["detections"], result["frame"]