def generate(self): # Get RGB-D image from camera image_bundle = self.camera.get_image_bundle() rgb = image_bundle['rgb'] depth = image_bundle['aligned_depth'] x, depth_img, rgb_img = self.cam_data.get_data(rgb=rgb, depth=depth) # Predict the grasp pose using the saved model with torch.no_grad(): xc = x.to(self.device) pred = self.model.predict(xc) q_img, ang_img, width_img = post_process_output( pred['pos'], pred['cos'], pred['sin'], pred['width']) grasps = detect_grasps(q_img, ang_img, width_img) # Get grasp position from model output pos_z = depth[grasps[0].center[0] + self.cam_data.top_left[0], grasps[0].center[1] + self.cam_data.top_left[1]] * self.cam_depth_scale - 0.04 pos_x = np.multiply( grasps[0].center[1] + self.cam_data.top_left[1] - self.camera.intrinsics.ppx, pos_z / self.camera.intrinsics.fx) pos_y = np.multiply( grasps[0].center[0] + self.cam_data.top_left[0] - self.camera.intrinsics.ppy, pos_z / self.camera.intrinsics.fy) if pos_z == 0: return target = np.asarray([pos_x, pos_y, pos_z]) target.shape = (3, 1) print('target: ', target) # Convert camera to robot coordinates camera2robot = self.cam_pose target_position = np.dot(camera2robot[0:3, 0:3], target) + camera2robot[0:3, 3:] target_position = target_position[0:3, 0] # Convert camera to robot angle angle = np.asarray([0, 0, grasps[0].angle]) angle.shape = (3, 1) target_angle = np.dot(camera2robot[0:3, 0:3], angle) # Concatenate grasp pose with grasp angle grasp_pose = np.append(target_position, target_angle[2]) print('grasp_pose: ', grasp_pose) np.save(self.grasp_pose, grasp_pose) if self.fig: plot_grasp(fig=self.fig, rgb_img=self.cam_data.get_rgb(rgb, False), grasps=grasps, save=True)
def validate(net, device, val_data, batches_per_epoch): """ Run validation. :param net: Network :param device: Torch device :param val_data: Validation Dataset :param batches_per_epoch: Number of batches to run :return: Successes, Failures and Losses """ net.eval() results = {'correct': 0, 'failed': 0, 'loss': 0, 'losses': {}} ld = len(val_data) with torch.no_grad(): batch_idx = 0 while batch_idx < batches_per_epoch: for x, y, didx, rot, zoom_factor in val_data: batch_idx += 1 if batches_per_epoch is not None and batch_idx >= batches_per_epoch: break xc = x.to(device) yc = [yy.to(device) for yy in y] lossd = net.compute_loss(xc, yc) loss = lossd['loss'] results['loss'] += loss.item() / ld for ln, l in lossd['losses'].items(): if ln not in results['losses']: results['losses'][ln] = 0 results['losses'][ln] += l.item() / ld q_out, ang_out, w_out = post_process_output( lossd['pred']['pos'], lossd['pred']['cos'], lossd['pred']['sin'], lossd['pred']['width']) s = evaluation.calculate_iou_match( q_out, ang_out, val_data.dataset.get_gtbb(didx, rot, zoom_factor), no_grasps=1, grasp_width=w_out, ) if s: results['correct'] += 1 else: results['failed'] += 1 return results
def grasp(): # Load image logging.info('Loading image...') pic = Image.open('rgb.png', 'r') rgb = np.array(pic) pic = Image.open('depth.png', 'r') depth = np.expand_dims(np.array(pic), axis=2) # Load Network logging.info('Loading model...') net = torch.load(args.network, map_location=torch.device('cpu')) logging.info('Done') # Get the compute device device = get_device(False) img_data = CameraData(include_depth=1, include_rgb=1) x, depth_img, rgb_img = img_data.get_data(rgb=rgb, depth=depth) with torch.no_grad(): xc = x.to(device) pred = net.predict(xc) q_img, ang_img, width_img = post_process_output( pred['pos'], pred['cos'], pred['sin'], pred['width']) if args.save: save_results(rgb_img=img_data.get_rgb(rgb, False), depth_img=np.squeeze(img_data.get_depth(depth)), grasp_q_img=q_img, grasp_angle_img=ang_img, no_grasps=args.n_grasps, grasp_width_img=width_img) else: fig = plt.figure(figsize=(12, 3)) gs = plot_results(fig=fig, rgb_img=img_data.get_rgb(rgb, False), grasp_q_img=q_img, grasp_angle_img=ang_img, no_grasps=args.n_grasps, grasp_width_img=width_img) fig.savefig('img_result.png')
def predict_grasp_angle(network, rgb_path, depth_path): #args = parse_args() #network = "trained-models/cornell-randsplit-rgbd-grconvnet3-drop1-ch16/epoch_30_iou_0.97" #rgb_path = "C:/Users/yashs/OneDrive/Desktop/PS simulation/rgbd_images/color8.jpeg" #depth_path = "C:/Users/yashs/OneDrive/Desktop/PS simulation/rgbd_images/depth8.jpeg" use_depth = 1 use_rgb = 1 n_grasps = 1 save = 0 force_cpu = False # Load image logging.info('Loading image...') pic = Image.open(rgb_path, 'r') rgb = np.array(pic) pic = Image.open(depth_path, 'r') depth = np.expand_dims(np.array(pic), axis=2) # Load Network logging.info('Loading model...') net = torch.load(network,map_location=torch.device('cpu')) logging.info('Done') # Get the compute device device = get_device(force_cpu) img_data = CameraData(include_depth=use_depth, include_rgb=use_rgb) x, depth_img, rgb_img = img_data.get_data(rgb=rgb, depth=depth) with torch.no_grad(): xc = x.to(device) pred = net.predict(xc) q_img, ang_img, width_img = post_process_output(pred['pos'], pred['cos'], pred['sin'], pred['width']) #print(pred['pos'].size()) #print(pred['pos']) #print(pred['cos']) #print(pred['sin']) #print(pred['width']) if save: save_results( rgb_img=img_data.get_rgb(rgb, False), depth_img=np.squeeze(img_data.get_depth(depth)), grasp_q_img=q_img, grasp_angle_img=ang_img, no_grasps=n_grasps, grasp_width_img=width_img ) else: fig = plt.figure(figsize=(10, 10)) gs=plot_results(fig=fig, rgb_img=img_data.get_rgb(rgb, False), grasp_q_img=q_img, grasp_angle_img=ang_img, no_grasps=n_grasps, grasp_width_img=width_img) fig.savefig('img_result.pdf') for g in gs: print(g.center) print(g.angle) print(g.length) print(g.width) return gs #predict_grasp_angle("trained-models/cornell-randsplit-rgbd-grconvnet3-drop1-ch16/epoch_notbest_17_iou_0.00", "C:/Users/yashs/OneDrive/Desktop/PS simulation/rgbd_images/color8.png", "C:/Users/yashs/OneDrive/Desktop/PS simulation/rgbd_images/depth8.png")
if args.jacquard_output: jo_fn = network + '_jacquard_output.txt' with open(jo_fn, 'w') as f: pass start_time = time.time() with torch.no_grad(): for idx, (x, y, didx, rot, zoom) in enumerate(test_data): xc = x.to(device) yc = [yi.to(device) for yi in y] lossd = net.compute_loss(xc, yc) q_img, ang_img, width_img = post_process_output( lossd['pred']['pos'], lossd['pred']['cos'], lossd['pred']['sin'], lossd['pred']['width']) if args.iou_eval: s = evaluation.calculate_iou_match( q_img, ang_img, test_data.dataset.get_gtbb(didx, rot, zoom), no_grasps=args.n_grasps, grasp_width=width_img, threshold=args.iou_threshold) if s: results['correct'] += 1 else: results['failed'] += 1
net = torch.load(args.network) logging.info('Done') # Get the compute device device = get_device(args.force_cpu) img_data = CameraData(include_depth=args.use_depth, include_rgb=args.use_rgb) x, depth_img, rgb_img = img_data.get_data(rgb=rgb, depth=depth) with torch.no_grad(): xc = x.to(device) pred = net.predict(xc) q_img, ang_img, width_img = post_process_output( pred['pos'], pred['cos'], pred['sin'], pred['width']) if args.save: save_results(rgb_img=img_data.get_rgb(rgb, False), depth_img=np.squeeze(img_data.get_depth(depth)), grasp_q_img=q_img, grasp_angle_img=ang_img, no_grasps=args.n_grasps, grasp_width_img=width_img) else: fig = plt.figure(figsize=(10, 10)) plot_results(fig=fig, rgb_img=img_data.get_rgb(rgb, False), grasp_q_img=q_img, grasp_angle_img=ang_img, no_grasps=args.n_grasps,
pass start_time = time.time() with torch.no_grad(): for idx, (x, y, didx, rot, zoom) in enumerate(test_data): xc = x.to(device) yc = [yi.to(device) for yi in y] # lossd = net.compute_loss(xc) # q_img, ang_img, width_img = post_process_output(lossd['pred']['pos'], lossd['pred']['cos'], # lossd['pred']['sin'], lossd['pred']['width']) pred = net(xc) q_img, ang_img, width_img = post_process_output(pred[0], pred[1], pred[2], pred[3]) if args.iou_eval: s = evaluation.calculate_iou_match(q_img, ang_img, test_data.dataset.get_gtbb(didx, rot, zoom), no_grasps=args.n_grasps, grasp_width=width_img, threshold=args.iou_threshold ) if s: results['correct'] += 1 else: results['failed'] += 1 if args.jacquard_output: grasps = grasp.detect_grasps(q_img, ang_img, width_img=width_img, no_grasps=1)