-
Notifications
You must be signed in to change notification settings - Fork 1
/
cube_callibration.py
92 lines (73 loc) · 2.76 KB
/
cube_callibration.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
# Author: Bharath Kumar
# Contact: kumar7bharath@gmail.com
import cv2
import numpy as np
from enum import Enum
# cube center points in world frame
# x - right, y - up, z - away
TILE_SIZE = 1.6
TILES_X = 4
TILES_Y = 4
CUBE_CORNER_WF = np.array([[(TILE_SIZE/2 + y*TILE_SIZE, TILE_SIZE/2 + x*TILE_SIZE) for x in range(TILES_X)] for y in range(TILES_Y)])
COLOR_LIMITS = np.array([[(79, 0, 0), (99, 255, 255)],
[(0, 155, 0), (17, 255, 255)],
[(112, 155, 0), (129, 255, 255)],
[(18, 0, 0), (43, 255, 255)],
[(0, 0, 0), (179, 20, 255)]])
class Colors(Enum):
GREEN = 4
ORANGE = 3
VIOLET = 2
YELLOW = 1
WHITE = 5
class CameraCallibration():
def __init__(self, image):
self.image = image
def get_tiles_centers(self):
result_image = self.image
image = self.separate_tile_colors()
ret, thresh = cv2.threshold(image, 150, 255, 0)
kernel = np.ones((5,5),np.uint8)
erosion = cv2.erode(thresh,kernel,iterations = 2)
contours, hierarchy = cv2.findContours(erosion, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours:
area = cv2.contourArea(contour)
M = cv2.moments(contour)
if M["m00"]:
cX = int(M["m10"] / M["m00"])
cY = int(M["m01"] / M["m00"])
if 100 < area < 100000:
cv2.circle(result_image, (cX, cY), 5, [255, 0, 0], -1)
cv2.imshow("frame", result_image)
cv2.waitKey(0)
def separate_tile_colors(self):
image = self.image
img_hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
result = np.zeros(image.shape[:2], dtype="uint8")
for color in range(len(COLOR_LIMITS)):
low_thresh = np.array(COLOR_LIMITS[color][0])
high_thresh = np.array(COLOR_LIMITS[color][1])
mask = cv2.inRange(img_hsv, low_thresh, high_thresh)
temp = cv2.GaussianBlur(mask, (25, 25), 1)
temp = cv2.erode(temp, (25, 25))
result = cv2.bitwise_or(result, temp)
# cv2.imshow("frame", result)
# cv2.waitKey(0)
return result
def test_video():
vid = cv2.VideoCapture("data/sample_video1.mp4")
print("inside")
while True:
ret, frame = vid.read()
callib = CameraCallibration(frame)
erosion = callib.get_tiles_centers()
cv2.imshow("output", erosion)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
vid.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
image = cv2.imread("callib_data/1.jpg")
callib = CameraCallibration(image)
callib.get_tiles_centers()
# callib.separate_tile_colors()