def __init__(self, hyperparams): # Get configs, hyperparameters & initialize agent self.ptconf, self.dconf, self.camconf = hyperparams[ 'ptconf'], hyperparams['dconf'], hyperparams['camconf'] self.env_name = hyperparams['env'] self.gui_on = hyperparams['bullet_gui_on'] self.hyperparams = hyperparams config = copy(AGENT_BULLET) config.update(hyperparams) Agent.__init__(self, config) # Setup bullet environment self._setup_conditions() self._setup_world(hyperparams['filename']) self.setup_bullet() self.setup_inference_camera() # Get demo data self.demo_vid = imageio.get_reader( join(self.dconf.DEMO_DIR, self.dconf.DEMO_NAME, 'rgb/{}.mp4'.format(self.dconf.SEQNAME))) self.demo_frames = [] self.reset_condition = hyperparams['reset_condition'] for im in self.demo_vid: self.demo_frames.append(im) # self.cost_tgt_mean = hyperparams['cost_tgt_mean'] # self.cost_tgt_std = hyperparams['cost_tgt_std'] self.objects_centroids_mrcnn = hyperparams['objects_centroids_mrcnn'] # Setup feature embedding network if enabled self.tcn = None # Setup Mask RCNN inference_config = InferenceConfig() with tf.device('/device:GPU:1'): self.mrcnn = modellib.MaskRCNN( mode='inference', model_dir=join(self.ptconf.EXP_DIR, self.ptconf.EXP_NAME, 'mrcnn_logs'), config=inference_config) self.mrcnn.load_weights(self.ptconf.WEIGHTS_FILE_PATH_MRCNN, by_name=True) self.class_names = gconf.CLASS_NAMES_W_BG # self.target_ids = gconf.CLASS_IDS self.target_ids = [1, 2] self.colors = visualize.random_colors(7) self.plotting_on = hyperparams['plotting_on'] if self.plotting_on: self.fig, self.ax = visualize.get_ax() self.mrcnn_centroids_last_known = { key: None for key in self.target_ids } st() # Run MRCNN on one image because first detection always takes long to initialize self.vid_seqname = 0 rgb_crop, depth_crop = self.get_images(0) results = self.mrcnn.detect([rgb_crop], verbose=0) self.blob_detector = BlobDetector()
def _get_rcnn_features(self, image, depth_rescaled): results = self.rcnn.detect([image], verbose=0) r = results[0] encountered_ids = [] all_cropped_boxes = [] all_centroids_unordered = [] # X Y Z all_centroids = dict() all_visual_features_unordered = [] all_visual_features = dict() for i, box in enumerate(r['rois']): class_id = r['class_ids'][i] if class_id not in self.target_ids or class_id in encountered_ids: continue encountered_ids.append(class_id) cropped = utils.crop_box(image, box, y_offset=20, x_offset=20) # cropped = utils.resize_image(cropped, max_dim=299)[0] cropped = cv2.resize(cropped, (299, 299)) all_cropped_boxes.append(cropped) masked_depth = depth_rescaled * r['masks'][:, :, i] masked_depth = masked_depth[np.where(masked_depth > 0)] z = np.median(np.sort(masked_depth.flatten())) x, y = utils.get_box_center(box) all_centroids_unordered.append([x, y, z]) all_visual_features_unordered.append(r['roi_features'][i]) all_cropped_boxes = np.asarray(all_cropped_boxes) all_centroids_unordered = np.asarray(all_centroids_unordered) for i in range(all_cropped_boxes.shape[0]): all_visual_features[ encountered_ids[i]] = all_visual_features_unordered[i] all_centroids[encountered_ids[i]] = all_centroids_unordered[i] all_centroids = np.asarray([val for key, val in all_centroids.items()]) all_visual_features = np.asarray( [val for key, val in all_visual_features.items()]) if self.plot_mode: fig, ax = visualize.get_ax() ax = visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], self.class_names, r['scores'], ax=ax, colors=self.colors) else: fig = None return all_visual_features, all_centroids, fig
def get_raw_rcnn_results(self, image): results = self.rcnn.detect([image], verbose=0) r = results[0] if self.plot_mode: fig, ax = visualize.get_ax() ax = visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], self.class_names, r['scores'], ax=ax, colors=self.colors) else: fig = None return r, fig
step_offset_back:step_size] imgs = imgs[TRUNCATE_FRONT + step_offset_front:-TRUNCATE_BACK - pose_offset - step_offset_back:step_size] dimgs = dimgs[TRUNCATE_FRONT + step_offset_front:-TRUNCATE_BACK - pose_offset - step_offset_back:step_size] print('length of truncated demo video: ', len(imgs)) ##### PLOTTING & COST PARSING ##### ### Compute pixel feature target and plot images ### tgt_imgs = [] plot_imgs = [] # target imgs stacked horizontally for plotting purposes points_tgt = [] if agent['visualize_mrcnn']: fig, ax = visualize.get_ax() else: ax = None # reader_w_robot = imageio.get_reader('/media/zhouxian/ed854110-6801-4dcd-9acf-c4f904955d71/iccv2019/gps_data/pour/lacan9_view0/data_files/2019-03-29_01-10-44/vids/rgb_sample_itr_9.mp4') # imgs_r = [] # for img in reader_w_robot: # imgs_r.append(img) # Compute MRCNN centroids & masks if not loaded from file # Compute MRCNN centroids & masks if not loaded from file for tt in range(agent['T']): dimg = dimgs[tt] img = imgs[tt] agent['demo_imgs'].append(img)
if FROM_DATASET: # # Test on a random image # Validation dataset image_ids = dataset.image_ids APs = [] for image_id in image_ids: image = dataset.load_image(image_id) results = model.detect([image], verbose=1) r = results[0] visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], dataset.class_names, r['scores'], ax=visualize.get_ax()[1], colors=colors) else: for seq_name in os.listdir('/home/msieb/projects/Mask_RCNN/datasets/baxter/test'): print("Processing ", seq_name) # seq_name = args.seqname dataset_dir = os.path.join(DATASET_DIR, "test", seq_name) filenames = os.listdir(dataset_dir) filenames = [file for file in filenames if '.jpg' in file] filenames = sorted(filenames, key=lambda x: x.split('.')[0]) for ii, file in enumerate(filenames): if not ii % 1 == 0: continue # # Load image and ground truth data # image, image_meta, gt_class_id, gt_bbox, gt_mask =\