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)
