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