def get_net():
# Init the model
net = GraspNet(input_feature_dim=0, num_view=cfgs.num_view, num_angle=12, num_depth=4,
cylinder_radius=0.05, hmin=-0.02, hmax_list=[0.01, 0.02, 0.03, 0.04], is_training=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# Load checkpoint
checkpoint = torch.load(cfgs.checkpoint_path)
net.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
print("-> loaded checkpoint %s (epoch: %d)" % (cfgs.checkpoint_path, start_epoch))
# set model to eval mode
net.eval()
return net
def __init__(self,
dataset_root,
dump_dir,
camera,
quick=True,
checkpoint_path=None):
self.dataset_root = dataset_root
self.dump_dir = dump_dir
self.camera = camera
self.quick = quick
if not self.quick:
self.collision_thresh = 0.01
self.voxel_size = 0.01
self.dataset = GraspNetDataset(dataset_root,
valid_obj_idxs=None,
grasp_labels=None,
split='all',
camera=camera,
remove_outlier=True,
augment=False,
load_label=False)
self.net = GraspNet(input_feature_dim=0,
num_view=300,
num_angle=12,
num_depth=4,
cylinder_radius=0.05,
hmin=-0.02,
hmax_list=[0.01, 0.02, 0.03, 0.04],
is_training=False)
self.net.to(device)
# Load checkpoint
assert checkpoint_path != None
checkpoint = torch.load(checkpoint_path)
self.net.load_state_dict(checkpoint['model_state_dict'])
print("-> loaded checkpoint %s." % (checkpoint_path))
self.net.eval()
shuffle=True,
num_workers=4,
worker_init_fn=my_worker_init_fn,
collate_fn=collate_fn)
TEST_DATALOADER = DataLoader(TEST_DATASET,
batch_size=cfgs.batch_size,
shuffle=False,
num_workers=4,
worker_init_fn=my_worker_init_fn,
collate_fn=collate_fn)
print(len(TRAIN_DATALOADER), len(TEST_DATALOADER))
# Init the model and optimzier
net = GraspNet(input_feature_dim=0,
num_view=cfgs.num_view,
num_angle=12,
num_depth=4,
cylinder_radius=0.05,
hmin=-0.02,
hmax_list=[0.01, 0.02, 0.03, 0.04])
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# Load the Adam optimizer
optimizer = optim.Adam(net.parameters(),
lr=cfgs.learning_rate,
weight_decay=cfgs.weight_decay)
# Load checkpoint if there is any
it = -1 # for the initialize value of `LambdaLR` and `BNMomentumScheduler`
start_epoch = 0
if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH):
checkpoint = torch.load(CHECKPOINT_PATH)
net.load_state_dict(checkpoint['model_state_dict'])
print('test dataset length: ', len(TEST_DATASET))
SCENE_LIST = TEST_DATASET.scene_list()
TEST_DATALOADER = DataLoader(TEST_DATASET,
batch_size=cfgs.batch_size,
shuffle=False,
num_workers=2,
worker_init_fn=my_worker_init_fn,
collate_fn=collate_fn)
print('test dataloader length: ', len(TEST_DATALOADER))
# Init the model
net = GraspNet(input_feature_dim=0,
num_view=cfgs.num_view,
num_angle=12,
num_depth=4,
cylinder_radius=0.05,
hmin=-0.02,
hmax_list=[0.01, 0.02, 0.03, 0.04],
is_training=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# Load checkpoint
checkpoint = torch.load(cfgs.checkpoint_path)
net.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
print("-> loaded checkpoint %s (epoch: %d)" %
(cfgs.checkpoint_path, start_epoch))
# ------------------------------------------------------------------------- GLOBAL CONFIG END
class GraspDetector():
def __init__(self,
dataset_root,
dump_dir,
camera,
quick=True,
checkpoint_path=None):
self.dataset_root = dataset_root
self.dump_dir = dump_dir
self.camera = camera
self.quick = quick
if not self.quick:
self.collision_thresh = 0.01
self.voxel_size = 0.01
self.dataset = GraspNetDataset(dataset_root,
valid_obj_idxs=None,
grasp_labels=None,
split='all',
camera=camera,
remove_outlier=True,
augment=False,
load_label=False)
self.net = GraspNet(input_feature_dim=0,
num_view=300,
num_angle=12,
num_depth=4,
cylinder_radius=0.05,
hmin=-0.02,
hmax_list=[0.01, 0.02, 0.03, 0.04],
is_training=False)
self.net.to(device)
# Load checkpoint
assert checkpoint_path != None
checkpoint = torch.load(checkpoint_path)
self.net.load_state_dict(checkpoint['model_state_dict'])
print("-> loaded checkpoint %s." % (checkpoint_path))
self.net.eval()
def get_grasp_group(self, scene_id, ann_id):
# print(ann_id)
# import ipdb; ipdb.set_trace()
camera_poses_path = os.path.join(self.dataset_root, 'scenes',
get_scene_name(scene_id), self.camera,
'camera_poses.npy')
if self.quick: # use cached grasp files
npy_file_path = os.path.join(self.dump_dir,
get_scene_name(scene_id), self.camera,
'%04d.npy' % (ann_id, ))
gg = AroundViewGraspGroup().from_npy(npy_file_path,
camera_poses_path)
else:
data_idx = scene_id * 256 + ann_id
data = self.dataset.__getitem__(data_idx)
data['point_clouds'] = to_var(data['point_clouds']).unsqueeze(0)
data['cloud_colors'] = to_var(data['cloud_colors']).unsqueeze(0)
# Forward pass
with torch.no_grad():
end_points = self.net(data)
grasp_preds = pred_decode(end_points)
preds = grasp_preds[0].detach().cpu().numpy()
gg = GraspGroup(preds)
# collision detection
if self.collision_thresh > 0:
cloud, _ = self.dataset.get_data(data_idx,
return_raw_cloud=True)
mfcdetector = ModelFreeCollisionDetector(
cloud, voxel_size=self.voxel_size)
collision_mask = mfcdetector.detect(
gg,
approach_dist=0.05,
collision_thresh=self.collision_thresh)
gg = gg[~collision_mask]
gg = AroundViewGraspGroup(np.array([ann_id] * len(gg)),
np.load(camera_poses_path),
gg.grasp_group_array)
return gg
def views2grasps(self, scene_id, views, offset):
if scene_id >= 4000: offset = 0
scene_id = get_sceneId4det(scene_id)
ann_ids = ALL_ANN_IDs[views] + offset
grasps = list()
for ann_id in ann_ids:
g = self.get_grasp_group(scene_id, ann_id)
grasps.append(g)
return grasps