-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathHandrailLocator.py
165 lines (126 loc) · 6.23 KB
/
HandrailLocator.py
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import cv2
import numpy as np
from Detector import Detector
from CameraController import CameraController
class CalibrationPackage:
def __init__(self, img_width, img_height, dimension_distance_list=[]):
self.image_width = img_height
self.image_height = img_height
self.dimension_distance_list = dimension_distance_list # width, height, distance
def add_detection(self, rotatedRect, dist):
dims = rotatedRect[1]
width = max(dims)
height = min(dims)
self.dimension_distance_list.append((width, height, dist))
return self
class HandrailFilter:
def __init__(self):
self.known_height = 3.50266 # cm
self.known_width = 25.5 # cm
self.focal_length_width = -1
self.focal_length_height = -1
self.image_width = 0
self.image_height = 0
self.detector = Detector()
def calibrate_from_package(self, calibration_pckg):
# [self.calibrate_from_detection(rect, dist) for rect, dist in calibration_pckg.detection_distances]
focal_lengths_list = [(self.get_focal_length_from_detection_width(width, dist),
self.get_focal_length_from_detection_height(height, dist))
for width, height, dist in calibration_pckg.dimension_distance_list]
self.focal_length_width = sum(pair[0] for pair in focal_lengths_list) / len(focal_lengths_list)
self.focal_length_height = sum(pair[1] for pair in focal_lengths_list) / len(focal_lengths_list)
self.image_height = calibration_pckg.image_height
self.image_width = calibration_pckg.image_width
def calibrate_from_detection(self, rotatedRect, known_distance):
dims = rotatedRect[1]
width = max(dims)
height = min(dims)
self.focal_length_width = self.get_focal_length_from_detection_width(width, known_distance)
self.focal_length_height = self.get_focal_length_from_detection_height(height, known_distance)
def get_focal_length_from_detection_height(self, height, known_distance):
return height * known_distance / self.known_height
def get_focal_length_from_detection_width(self, width, known_distance):
return width * known_distance / self.known_width
def get_distance_from_detection(self, rotatedRect):
dims = rotatedRect[1]
width = max(dims)
height = min(dims)
distance_h = self.get_distance_from_detection_height(height)
distance_w = self.get_distance_from_detection_width(width)
return min(distance_h, distance_w)
def get_distance_from_detection_height(self, height):
return self.known_height * self.focal_length_height / height
def get_distance_from_detection_width(self, width):
return self.known_width * self.focal_length_width / width
def is_handrail(self, rotatedRect):
dims = rotatedRect[1]
width = max(dims)
height = min(dims)
return width / height >= 3 # placeholder comparison
def remove_non_handrail_detections(self, rotatedRects):
filtered_rects = [rect for rect in rotatedRects if self.is_handrail(rect)]
return filtered_rects
def get_valid_detections_and_distances_list(self, rotatedRects):
return [(rect, self.get_distance_from_detection(rect))
for rect in self.remove_non_handrail_detections(rotatedRects)]
def get_vector_to_handrail(self, detection, distance):
pixel_to_m_conv = self.focal_length_width/distance
x_coord = detection[0][0]
x_offset_cm = abs(self.image_width / 2 - x_coord) / pixel_to_m_conv
return distance, x_offset_cm
# This is the function that will be called by the MainController
def get_handrail_vectors(self, image):
detections = self.detector.get_rects_from_bgr(image)
detection_distance_pairs = self.get_valid_detections_and_distances_list(detections)
handrail_vector_list = []
for detection, distance in detection_distance_pairs:
handrail_vector_list += self.get_vector_to_handrail(detection, distance)
return handrail_vector_list
def test_distance_calculator():
detector = Detector()
handrail_filter = calibrate_from_package_return_handrail_filter()
test_img = cv2.imread("Calibration_Images/cal_80cm_45degrees.jpg")
detection = detector.get_rects_from_bgr(test_img)[0]
print(handrail_filter.get_distance_from_detection(detection))
return handrail_filter.get_distance_from_detection(detection)
# create a test with multiple handrails in frame print(angle)
def test_cam_dist_calculator():
detector = Detector()
handrail_filter = calibrate_from_package_return_handrail_filter()
cam = CameraController()
test_img = cam.get_image()
cv2.imshow("before", test_img)
cv2.waitKey(0)
detections = detector.get_rects_from_bgr(test_img)
out = handrail_filter.get_valid_detections_and_distances_list(detections)
def test_cam_angle_calculator():
detector = Detector()
handrail_filter = calibrate_from_package_return_handrail_filter()
cam = CameraController()
for i in 1, 2:
test_img = cam.get_image()
out = handrail_filter.get_handrail_vectors(test_img)
print(out)
def test_package_calibration():
detector = Detector()
handrail_filter = calibrate_from_package_return_handrail_filter()
test_img = cv2.imread("Calibration_Images/cal_80cm_45degrees.jpg")
detection = detector.get_rects_from_bgr(test_img)[0]
print(handrail_filter.get_distance_from_detection(detection))
def calibrate_from_package_return_handrail_filter():
detector = Detector()
handrail_filter = HandrailFilter()
cal1 = cv2.imread("Calibration_Images/cal_1m_0degrees.jpg")
height, width, channels = cal1.shape
cal_pckg = CalibrationPackage(width, height)
det1 = detector.get_rects_from_bgr(cal1)[0]
cal2 = cv2.imread("Calibration_Images/cal_80cm_0degrees.jpg")
det2 = detector.get_rects_from_bgr(cal1)[0]
cal_pckg.add_detection(det1, 100).add_detection(det2, 80)
handrail_filter.calibrate_from_package(cal_pckg)
return handrail_filter
if __name__ == "__main__":
# test_cam_dist_calculator()
# test_distance_calculator()
# test_package_calibration()
test_cam_angle_calculator()