def on_sensor_frame(self, data):
rgb_image = Image.open(BytesIO(base64.b64decode(data['rgb_image'])))
rgb_image = np.asarray(rgb_image)
if rgb_image.shape != (256, 256, 3):
print('image shape not 256, 256, 3')
return None
if rgb_image.shape[0] != self.image_hw:
rgb_image = misc.imresize(rgb_image, (self.image_hw, self.image_hw, 3))
rgb_image = data_iterator.preprocess_input(rgb_image)
pred = np.squeeze(model.predict(np.expand_dims(rgb_image, 0)))
if self.pred_viz_enabled:
self.queue.put([rgb_image, pred])
target_mask = pred[:, :, 2] > 0.5
# reduce the number of false positives by requiring more pixels to be identified as containing the target
if target_mask.sum() > 10:
centroid = scoring_utils.get_centroid_largest_blob(target_mask)
# scale the centroid from the nn image size to the original image size
centroid = centroid.astype(np.int).tolist()
# Obtain 3D world point from centroid pixel
depth_img = get_depth_image(data['depth_image'])
# Get XYZ coordinates for specific pixel
pixel_depth = depth_img[centroid[0]][centroid[1]][0]*50/255.0
point_3d = get_xyz_from_image(centroid[0], centroid[1], pixel_depth, self.image_hw)
point_3d.append(1)
# Get cam_pose from sensor_frame (ROS convention)
cam_pose = get_ros_pose(data['gimbal_pose'])
# Calculate xyz-world coordinates of the point corresponding to the pixel
# Transformation Matrix
R = euler2mat(math.radians(cam_pose[3]), math.radians(cam_pose[4]), math.radians(cam_pose[5]))
T = np.c_[R, cam_pose[:3]]
T = np.vstack([T, [0,0,0,1]]) # transformation matrix from world to quad
# 3D point in ROS coordinates
ros_point = np.dot(T, point_3d)
sio.emit('object_detected', {'coords': [ros_point[0], ros_point[1], ros_point[2]]})
if not self.target_found:
print('Target found!')
self.target_found = True
self.num_no_see = 0
# Publish Hero Marker
marker_pos = [ros_point[0],ros_point[1], ros_point[2]] + [0, 0, 0]
marker_msg = sio_msgs.create_box_marker_msg(np.random.randint(99999), marker_pos)
sio.emit('create_box_marker', marker_msg)
elif self.target_found:
self.num_no_see += 1
# print(self.num_no_see)
if self.target_found and self.num_no_see > 6:
print('Target lost!')
sio.emit('object_lost', {'data': ''})
self.target_found = False
self.num_no_see = 0
os.makedirs(path)
if __name__ == '__main__':
if len(sys.argv) != 3:
raise ValueError(
'predict.py requires a model file name, and output folder name cli input'
)
model_file = sys.argv[1]
output_folder = sys.argv[2]
output_path = os.path.join('..', 'data', 'runs', output_folder)
make_dir_if_not_exist(output_path)
model = make_model.make_example_model()
model.load_weights(os.path.join('..', 'data', 'weights', model_file))
data_folder = os.path.join('..', 'data', 'validation')
file_names = sorted(
glob.glob(os.path.join(data_folder, 'images', '*.jpeg')))
for name in file_names:
image = misc.imread(name)
image = data_iterator.preprocess_input(image.astype(np.float32))
pred = model.predict_on_batch(np.expand_dims(image, 0))
base_name = os.path.basename(name).split('.')[0]
base_name = base_name + '_prediction.png'
misc.imsave(os.path.join(output_path, base_name),
np.squeeze((pred * 255).astype(np.uint8)))
def on_sensor_data(self, data):
rgb_image = Image.open(BytesIO(base64.b64decode(data['rgb_image'])))
rgb_image = np.asarray(rgb_image)
if rgb_image.shape != (256, 256, 3):
print('image shape not 256, 256, 3')
return None
rgb_image = data_iterator.preprocess_input(rgb_image)
pred = np.squeeze(model.predict(np.expand_dims(rgb_image, 0)))
target_mask = pred[:, :, 1] > 0.5
if self.pred_viz_enabled:
overlay_viz(rgb_image, pred)
# reduce the number of false positives by requiring more pixels to be identified as containing the target
if target_mask.sum() > 10:
# Temporary move so we only get the overlays with positive identification
if self.image_save_dir is not None:
if time.time() - self.last_time_saved > 1:
result = visualization.overlay_predictions(rgb_image, pred, None, 0.5, 1)
out_file = os.path.join(self.image_save_dir, 'overlay_' + str(time.time()) + '.png')
misc.imsave(out_file, result)
self.last_time_saved = time.time()
centroid = scoring_utils.get_centroid_largest_blob(target_mask)
# scale the centroid from the nn image size to the original image size
centroid = centroid.astype(np.int).tolist()
# Obtain 3D world point from centroid pixel
depth_img = get_depth_image(data['depth_image'])
# Get XYZ coordinates for specific pixel
pixel_depth = depth_img[centroid[0]][centroid[1]][0]*100/255.0
point_3d = get_xyz_from_image(centroid[0], centroid[1], pixel_depth)
point_3d.append(1)
#print("Pixel Depth: ", pixel_depth)
#print ("Hit point: ", point_3d)
# Get cam_pose from sensor_frame (ROS convention)
cam_pose = get_ros_pose(data['gimbal_pose'])
#print ("Quad Pose: ", cam_pose)
# Calculate xyz-world coordinates of the point corresponding to the pixel
# Transformation Matrix
R = euler2mat(math.radians(cam_pose[3]), math.radians(cam_pose[4]), math.radians(cam_pose[5]))
T = np.c_[R, cam_pose[:3]]
T = np.vstack([T, [0,0,0,1]]) # transformation matrix from world to quad
# 3D point in ROS coordinates
ros_point = np.dot(T, point_3d)
socketIO.emit('object_detected', {'coords': [ros_point[0], ros_point[1], ros_point[2]]})
self.target_found = True
self.num_no_see = 0
# Publish Hero Marker
marker_pos = [ros_point[0],ros_point[1], ros_point[2]] + [0, 0, 0]
marker_msg = sio_msgs.create_box_marker_msg(np.random.randint(99999), marker_pos)
socketIO.emit('create_box_marker', marker_msg)
# 3D point in Unity coordinates
#unity_point = get_unity_pose_from_ros(ros_point)
#print ros_point.shape
# with the depth image, and centroid from prediction we can compute
# the x,y,z coordinates where the ray intersects an object
# ray = ray_casting.cast_ray(data, [centroid[1], centroid[0]])
# pose = np.array(list(map(float, data['gimbal_pose'].split(','))))
# TODO add rotation of the camera with respect to the gimbal
# create the rotation matrix to rotate the sensors frame of reference to the world frame
# rot = transformations.euler_matrix(to_radians(pose[3]),
# to_radians(pose[4]),
# to_radians(pose[5]))[:3, :3]
# rotate array
# ray = np.dot(rot, np.array(ray))
elif self.target_found:
self.num_no_see += 1
if self.target_found and self.num_no_see > 8:
socketIO.emit('object_lost', {'data': ''})
self.target_found = False
self.num_no_see = 0
def on_sensor_frame(self, data):
rgb_image = Image.open(BytesIO(base64.b64decode(data['rgb_image'])))
rgb_image = np.asarray(rgb_image)
if rgb_image.shape != (256, 256, 3):
print('image shape not 256, 256, 3')
return None
if rgb_image.shape[0] != self.image_hw:
rgb_image = misc.imresize(rgb_image, (self.image_hw, self.image_hw, 3))
rgb_image = data_iterator.preprocess_input(rgb_image)
pred = np.squeeze(model.predict(np.expand_dims(rgb_image, 0)))
if self.pred_viz_enabled:
self.queue.put([rgb_image, pred])
target_mask = pred[:, :, 2] > 0.5
# reduce the number of false positives by requiring more pixels to be identified as containing the target
if target_mask.sum() > 10:
centroid = scoring_utils.get_centroid_largest_blob(target_mask)
# scale the centroid from the nn image size to the original image size
centroid = centroid.astype(np.int).tolist()
# Obtain 3D world point from centroid pixel
depth_img = get_depth_image(data['depth_image'])
# Get XYZ coordinates for specific pixel
pixel_depth = depth_img[centroid[0]][centroid[1]][0]*50/255.0
point_3d = get_xyz_from_image(centroid[0], centroid[1], pixel_depth, self.image_hw)
point_3d.append(1)
# Get cam_pose from sensor_frame (ROS convention)
cam_pose = get_ros_pose(data['gimbal_pose'])
# Calculate xyz-world coordinates of the point corresponding to the pixel
# Transformation Matrix
R = euler2mat(math.radians(cam_pose[3]), math.radians(cam_pose[4]), math.radians(cam_pose[5]))
T = np.c_[R, cam_pose[:3]]
T = np.vstack([T, [0,0,0,1]]) # transformation matrix from world to quad
# 3D point in ROS coordinates
ros_point = np.dot(T, point_3d)
sio.emit('object_detected', {'coords': [ros_point[0], ros_point[1], ros_point[2]]})
if not self.target_found:
print('Target found!')
self.target_found = True
self.num_no_see = 0
# Publish Hero Marker
marker_pos = [ros_point[0],ros_point[1], ros_point[2]] + [0, 0, 0]
marker_msg = sio_msgs.create_box_marker_msg(np.random.randint(99999), marker_pos)
sio.emit('create_box_marker', marker_msg)
elif self.target_found:
self.num_no_see += 1
# print(self.num_no_see)
if self.target_found and self.num_no_see > 6:
print('Target lost!')
sio.emit('object_lost', {'data': ''})
self.target_found = False
self.num_no_see = 0