-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathperception.py
167 lines (136 loc) · 6.9 KB
/
perception.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
import numpy as np
import cv2
# Identify pixels above the threshold
# Threshold of RGB > 160 does a nice job of identifying ground pixels only
def color_thresh(img, rgb_thresh=(160, 160, 160)):
# Create an array of zeros same xy size as img, but single channel
color_select = np.zeros_like(img[:,:,0])
# Require that each pixel be above all three threshold values in RGB
# above_thresh will now contain a boolean array with "True"
# where threshold was met
above_thresh = (img[:,:,0] > rgb_thresh[0]) \
& (img[:,:,1] > rgb_thresh[1]) \
& (img[:,:,2] > rgb_thresh[2])
# Index the array of zeros with the boolean array and set to 1
color_select[above_thresh] = 1
# Return the binary image
return color_select
def rock_thresh(img):
# Convert BGR to HSV
hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV, 3)
# Define range of yellow colors in HSV
lower_yellow = np.array([20, 150, 100], dtype='uint8')
upper_yellow = np.array([50, 255, 255], dtype='uint8')
# Threshold the HSV image to get only yellow colors
mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
return mask
# Define a function to convert from image coords to rover coords
def rover_coords(binary_img):
# Identify nonzero pixels
ypos, xpos = binary_img.nonzero()
# Calculate pixel positions with reference to the rover position being at the
# center bottom of the image.
x_pixel = -(ypos - binary_img.shape[0]).astype(np.float)
y_pixel = -(xpos - binary_img.shape[1]/2 ).astype(np.float)
return x_pixel, y_pixel
# Define a function to convert to radial coords in rover space
def to_polar_coords(x_pixel, y_pixel):
# Convert (x_pixel, y_pixel) to (distance, angle)
# in polar coordinates in rover space
# Calculate distance to each pixel
dist = np.sqrt(x_pixel**2 + y_pixel**2)
# Calculate angle away from vertical for each pixel
angles = np.arctan2(y_pixel, x_pixel)
return dist, angles
# Define a function to map rover space pixels to world space
def rotate_pix(xpix, ypix, yaw):
# Convert yaw to radians
yaw_rad = yaw * np.pi / 180
xpix_rotated = (xpix * np.cos(yaw_rad)) - (ypix * np.sin(yaw_rad))
ypix_rotated = (xpix * np.sin(yaw_rad)) + (ypix * np.cos(yaw_rad))
# Return the result
return xpix_rotated, ypix_rotated
def translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale):
# Apply a scaling and a translation
xpix_translated = (xpix_rot / scale) + xpos
ypix_translated = (ypix_rot / scale) + ypos
# Return the result
return xpix_translated, ypix_translated
# Define a function to apply rotation and translation (and clipping)
# Once you define the two functions above this function should work
def pix_to_world(xpix, ypix, xpos, ypos, yaw, world_size, scale):
# Apply rotation
xpix_rot, ypix_rot = rotate_pix(xpix, ypix, yaw)
# Apply translation
xpix_tran, ypix_tran = translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale)
# Perform rotation, translation and clipping all at once
x_pix_world = np.clip(np.int_(xpix_tran), 0, world_size - 1)
y_pix_world = np.clip(np.int_(ypix_tran), 0, world_size - 1)
# Return the result
return x_pix_world, y_pix_world
# Define a function to perform a perspective transform
def perspect_transform(img, src, dst):
M = cv2.getPerspectiveTransform(src, dst)
warped = cv2.warpPerspective(img, M, (img.shape[1], img.shape[0]))# keep same size as input image
mask = cv2.warpPerspective(np.ones_like(img[:,:,0]), M, (img.shape[1], img.shape[0]))
return warped, mask
# Apply the above functions in succession and update the Rover state accordingly
def perception_step(Rover):
# Perform perception steps to update Rover()
# TODO:
# NOTE: camera image is coming to you in Rover.img
# 1) Define source and destination points for perspective transform
dst_size = 5
bottom_offset = 6
image = Rover.img
source = np.float32([[14, 140], [301 ,140],[200, 96], [118, 96]])
destination = np.float32([[image.shape[1]/2 - dst_size, image.shape[0] - bottom_offset],
[image.shape[1]/2 + dst_size, image.shape[0] - bottom_offset],
[image.shape[1]/2 + dst_size, image.shape[0] - 2*dst_size - bottom_offset],
[image.shape[1]/2 - dst_size, image.shape[0] - 2*dst_size - bottom_offset],
])
# 2) Apply perspective transform
warped, mask = perspect_transform(image, source, destination)
# 3) Apply color threshold to identify navigable terrain/obstacles/rock samples
color_select = color_thresh(warped, rgb_thresh=(160,160,160))
obstacle = np.absolute(np.float32(color_select) -1) * mask
rock = rock_thresh(warped)
# 4) Update Rover.vision_image (this will be displayed on left side of screen)
# Example: Rover.vision_image[:,:,0] = obstacle color-thresholded binary image
# Rover.vision_image[:,:,1] = rock_sample color-thresholded binary image
# Rover.vision_image[:,:,2] = navigable terrain color-thresholded binary image
Rover.vision_image[:,:,0] = obstacle * 255
Rover.vision_image[:,:,1] = rock * 255
Rover.vision_image[:,:,2] = color_select * 255
# 5) Convert map image pixel values to rover-centric coords
x_pix, y_pix = rover_coords(color_select)
obsxpix, obsypix = rover_coords(obstacle)
rockxpix, rockypix = rover_coords(rock)
# 6) Convert rover-centric pixel values to world coordinates
world_size = Rover.worldmap.shape[0]
scale = 2 * dst_size
xpos, ypos = Rover.pos[0], Rover.pos[1]
yaw = Rover.yaw
xpix_world, ypix_world = pix_to_world(x_pix, y_pix, xpos, ypos, yaw, world_size, scale)
rockxpix_world, rockypix_world = pix_to_world(rockxpix, rockypix, xpos, ypos, yaw, world_size, scale)
obsxpix_world, obsypix_world = pix_to_world(obsxpix, obsypix, xpos, ypos, yaw, world_size, scale)
# 7) Update Rover worldmap (to be displayed on right side of screen)
# Example: Rover.worldmap[obstacle_y_world, obstacle_x_world, 0] += 1
# Rover.worldmap[rock_y_world, rock_x_world, 1] += 1
# Rover.worldmap[navigable_y_world, navigable_x_world, 2] += 1
Rover.worldmap[obsypix_world, obsxpix_world, 0] += 255
Rover.worldmap[rockypix_world, rockxpix_world, 1] += 255
Rover.worldmap[ypix_world, xpix_world, 2] += 255
if len(rockxpix) > 5:
rock_dist, rock_angles = to_polar_coords(rockxpix, rockypix)
Rover.sample_angles = rock_angles
Rover.sample_dists = rock_dist
Rover.sample_seen = True
# 8) Convert rover-centric pixel positions to polar coordinates
dist, angles = to_polar_coords(x_pix, y_pix)
# Update Rover pixel distances and angles
# Rover.nav_dists = rover_centric_pixel_distances
# Rover.nav_angles = rover_centric_angles
Rover.nav_dists = dist
Rover.nav_angles = angles
return Rover