lokesh341 commited on
Commit
96d750e
·
1 Parent(s): 5a67b15

Update services/operations_maintenance/signage_check.py

Browse files
services/operations_maintenance/signage_check.py CHANGED
@@ -1,8 +1,6 @@
1
- # services/operations_maintenance/signage_check.py
2
  import cv2
3
  import numpy as np
4
  from ultralytics import YOLO
5
- import pytesseract
6
  import os
7
 
8
  # Load YOLOv8m model
@@ -10,22 +8,20 @@ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
10
  MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m.pt")
11
  model = YOLO(MODEL_PATH)
12
 
13
- def detect_signages(frame):
14
  """
15
- Detect signages in a frame using YOLOv8m and extract text with Tesseract.
16
  Args:
17
  frame: Input frame (numpy array)
18
  Returns:
19
- dict: Detection results with numbered labels
20
  numpy array: Annotated frame
21
  """
22
- # Run YOLOv8 inference
23
  results = model(frame)
24
 
25
  detections = []
26
  line_counter = 1 # Initialize counter for numbered labels
27
 
28
- # Process detections
29
  for r in results:
30
  for box in r.boxes:
31
  conf = float(box.conf[0])
@@ -38,32 +34,21 @@ def detect_signages(frame):
38
  xyxy = box.xyxy[0].cpu().numpy()
39
  x_min, y_min, x_max, y_max = map(int, xyxy)
40
 
41
- # Extract text from signage
42
- signage_roi = frame[y_min:y_max, x_min:x_max]
43
- text = pytesseract.image_to_string(signage_roi).strip()
44
-
45
  # Add numbered label
46
- detection_label = f"Line {line_counter} - Signage: {text} (Conf: {conf:.2f})"
47
  detections.append({
 
48
  "label": detection_label,
49
  "confidence": conf,
50
- "coordinates": [x_min, y_min, x_max, y_max],
51
- "text": text
52
  })
53
 
54
  # Draw bounding box and label
55
- color = (0, 255, 255) # Yellow for signages
56
  cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
57
  cv2.putText(frame, detection_label, (x_min, y_min - 10),
58
  cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
59
 
60
  line_counter += 1
61
 
62
- return {"detections": detections, "frame": frame}
63
-
64
- def process_signages(frame):
65
- """
66
- Wrapper function for integration with app.py.
67
- """
68
- result = detect_signages(frame)
69
- return result["detections"], result["frame"]
 
 
1
  import cv2
2
  import numpy as np
3
  from ultralytics import YOLO
 
4
  import os
5
 
6
  # Load YOLOv8m model
 
8
  MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m.pt")
9
  model = YOLO(MODEL_PATH)
10
 
11
+ def process_signages(frame):
12
  """
13
+ Detect road signage in a frame.
14
  Args:
15
  frame: Input frame (numpy array)
16
  Returns:
17
+ list: List of detected signage
18
  numpy array: Annotated frame
19
  """
 
20
  results = model(frame)
21
 
22
  detections = []
23
  line_counter = 1 # Initialize counter for numbered labels
24
 
 
25
  for r in results:
26
  for box in r.boxes:
27
  conf = float(box.conf[0])
 
34
  xyxy = box.xyxy[0].cpu().numpy()
35
  x_min, y_min, x_max, y_max = map(int, xyxy)
36
 
 
 
 
 
37
  # Add numbered label
38
+ detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
39
  detections.append({
40
+ "type": label,
41
  "label": detection_label,
42
  "confidence": conf,
43
+ "coordinates": [x_min, y_min, x_max, y_max]
 
44
  })
45
 
46
  # Draw bounding box and label
47
+ color = (0, 191, 255) # DeepSkyBlue for signage
48
  cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
49
  cv2.putText(frame, detection_label, (x_min, y_min - 10),
50
  cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
51
 
52
  line_counter += 1
53
 
54
+ return detections, frame