197 lines
8.3 KiB
Python
197 lines
8.3 KiB
Python
|
import os
|
||
|
|
import sys
|
||
|
|
import numpy as np
|
||
|
|
import argparse
|
||
|
|
from PIL import Image
|
||
|
|
import time
|
||
|
|
import scipy.io as scio
|
||
|
|
import torch
|
||
|
|
import open3d as o3d
|
||
|
|
from graspnetAPI.graspnet_eval import GraspGroup
|
||
|
|
|
||
|
|
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
|
|
sys.path.append(ROOT_DIR)
|
||
|
|
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
|
||
|
|
from models.graspnet import GraspNet, pred_decode
|
||
|
|
from dataset.graspnet_dataset import minkowski_collate_fn
|
||
|
|
from collision_detector import ModelFreeCollisionDetector
|
||
|
|
from data_utils import CameraInfo, create_point_cloud_from_depth_image, get_workspace_mask
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--dataset_root', default='/data/datasets/graspnet')
|
||
|
|
parser.add_argument('--checkpoint_path', default='/data/zibo/logs/graspness_kn.tar')
|
||
|
|
parser.add_argument('--dump_dir', help='Dump dir to save outputs', default='/data/zibo/logs/')
|
||
|
|
parser.add_argument('--seed_feat_dim', default=512, type=int, help='Point wise feature dim')
|
||
|
|
parser.add_argument('--camera', default='kinect', help='Camera split [realsense/kinect]')
|
||
|
|
parser.add_argument('--num_point', type=int, default=15000, help='Point Number [default: 15000]')
|
||
|
|
parser.add_argument('--batch_size', type=int, default=1, help='Batch Size during inference [default: 1]')
|
||
|
|
parser.add_argument('--voxel_size', type=float, default=0.005, help='Voxel Size for sparse convolution')
|
||
|
|
parser.add_argument('--collision_thresh', type=float, default=-1,
|
||
|
|
help='Collision Threshold in collision detection [default: 0.01]')
|
||
|
|
parser.add_argument('--voxel_size_cd', type=float, default=0.01, help='Voxel Size for collision detection')
|
||
|
|
parser.add_argument('--infer', action='store_true', default=False)
|
||
|
|
parser.add_argument('--vis', action='store_true', default=False)
|
||
|
|
parser.add_argument('--scene', type=str, default='0188')
|
||
|
|
parser.add_argument('--index', type=str, default='0000')
|
||
|
|
cfgs = parser.parse_args()
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------------- GLOBAL CONFIG BEG
|
||
|
|
if not os.path.exists(cfgs.dump_dir):
|
||
|
|
os.mkdir(cfgs.dump_dir)
|
||
|
|
|
||
|
|
|
||
|
|
def data_process():
|
||
|
|
root = cfgs.dataset_root
|
||
|
|
camera_type = cfgs.camera
|
||
|
|
|
||
|
|
depth = np.array(Image.open(os.path.join(root, 'scenes', scene_id, camera_type, 'depth', index + '.png')))
|
||
|
|
seg = np.array(Image.open(os.path.join(root, 'scenes', scene_id, camera_type, 'label', index + '.png')))
|
||
|
|
meta = scio.loadmat(os.path.join(root, 'scenes', scene_id, camera_type, 'meta', index + '.mat'))
|
||
|
|
try:
|
||
|
|
intrinsic = meta['intrinsic_matrix']
|
||
|
|
factor_depth = meta['factor_depth']
|
||
|
|
except Exception as e:
|
||
|
|
print(repr(e))
|
||
|
|
camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2],
|
||
|
|
factor_depth)
|
||
|
|
# generate cloud
|
||
|
|
cloud = create_point_cloud_from_depth_image(depth, camera, organized=True)
|
||
|
|
|
||
|
|
# get valid points
|
||
|
|
depth_mask = (depth > 0)
|
||
|
|
camera_poses = np.load(os.path.join(root, 'scenes', scene_id, camera_type, 'camera_poses.npy'))
|
||
|
|
align_mat = np.load(os.path.join(root, 'scenes', scene_id, camera_type, 'cam0_wrt_table.npy'))
|
||
|
|
trans = np.dot(align_mat, camera_poses[int(index)])
|
||
|
|
workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02)
|
||
|
|
mask = (depth_mask & workspace_mask)
|
||
|
|
|
||
|
|
cloud_masked = cloud[mask]
|
||
|
|
|
||
|
|
# sample points random
|
||
|
|
if len(cloud_masked) >= cfgs.num_point:
|
||
|
|
idxs = np.random.choice(len(cloud_masked), cfgs.num_point, replace=False)
|
||
|
|
else:
|
||
|
|
idxs1 = np.arange(len(cloud_masked))
|
||
|
|
idxs2 = np.random.choice(len(cloud_masked), cfgs.num_point - len(cloud_masked), replace=True)
|
||
|
|
idxs = np.concatenate([idxs1, idxs2], axis=0)
|
||
|
|
cloud_sampled = cloud_masked[idxs]
|
||
|
|
|
||
|
|
ret_dict = {'point_clouds': cloud_sampled.astype(np.float32),
|
||
|
|
'coors': cloud_sampled.astype(np.float32) / cfgs.voxel_size,
|
||
|
|
'feats': np.ones_like(cloud_sampled).astype(np.float32),
|
||
|
|
}
|
||
|
|
return ret_dict
|
||
|
|
|
||
|
|
|
||
|
|
# Init datasets and dataloaders
|
||
|
|
def my_worker_init_fn(worker_id):
|
||
|
|
np.random.seed(np.random.get_state()[1][0] + worker_id)
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def inference(data_input):
|
||
|
|
batch_data = minkowski_collate_fn([data_input])
|
||
|
|
net = GraspNet(seed_feat_dim=cfgs.seed_feat_dim, 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))
|
||
|
|
|
||
|
|
net.eval()
|
||
|
|
tic = time.time()
|
||
|
|
|
||
|
|
for key in batch_data:
|
||
|
|
if 'list' in key:
|
||
|
|
for i in range(len(batch_data[key])):
|
||
|
|
for j in range(len(batch_data[key][i])):
|
||
|
|
batch_data[key][i][j] = batch_data[key][i][j].to(device)
|
||
|
|
else:
|
||
|
|
batch_data[key] = batch_data[key].to(device)
|
||
|
|
# Forward pass
|
||
|
|
with torch.no_grad():
|
||
|
|
end_points = net(batch_data)
|
||
|
|
grasp_preds = pred_decode(end_points)
|
||
|
|
|
||
|
|
preds = grasp_preds[0].detach().cpu().numpy()
|
||
|
|
|
||
|
|
# Filtering grasp poses for real-world execution.
|
||
|
|
# The first mask preserves the grasp poses that are within a 30-degree angle with the vertical pose and have a width of less than 9cm.
|
||
|
|
# mask = (preds[:,9] > 0.9) & (preds[:,1] < 0.09)
|
||
|
|
# The second mask preserves the grasp poses within the workspace of the robot.
|
||
|
|
# workspace_mask = (preds[:,12] > -0.20) & (preds[:,12] < 0.21) & (preds[:,13] > -0.06) & (preds[:,13] < 0.18) & (preds[:,14] > 0.63)
|
||
|
|
# preds = preds[mask & workspace_mask]
|
||
|
|
|
||
|
|
# if len(preds) == 0:
|
||
|
|
# print('No grasp detected after masking')
|
||
|
|
# return
|
||
|
|
|
||
|
|
gg = GraspGroup(preds)
|
||
|
|
# collision detection
|
||
|
|
if cfgs.collision_thresh > 0:
|
||
|
|
cloud = data_input['point_clouds']
|
||
|
|
mfcdetector = ModelFreeCollisionDetector(cloud, voxel_size=cfgs.voxel_size_cd)
|
||
|
|
collision_mask = mfcdetector.detect(gg, approach_dist=0.05, collision_thresh=cfgs.collision_thresh)
|
||
|
|
gg = gg[~collision_mask]
|
||
|
|
|
||
|
|
# save grasps
|
||
|
|
save_dir = os.path.join(cfgs.dump_dir, scene_id, cfgs.camera)
|
||
|
|
save_path = os.path.join(save_dir, cfgs.index + '.npy')
|
||
|
|
if not os.path.exists(save_dir):
|
||
|
|
os.makedirs(save_dir)
|
||
|
|
gg.save_npy(save_path)
|
||
|
|
|
||
|
|
toc = time.time()
|
||
|
|
print('inference time: %fs' % (toc - tic))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
scene_id = 'scene_' + cfgs.scene
|
||
|
|
index = cfgs.index
|
||
|
|
data_dict = data_process()
|
||
|
|
|
||
|
|
if cfgs.infer:
|
||
|
|
inference(data_dict)
|
||
|
|
if cfgs.vis:
|
||
|
|
pc = data_dict['point_clouds']
|
||
|
|
gg = np.load(os.path.join(cfgs.dump_dir, scene_id, cfgs.camera, cfgs.index + '.npy'))
|
||
|
|
gg = GraspGroup(gg)
|
||
|
|
gg = gg.nms()
|
||
|
|
gg = gg.sort_by_score()
|
||
|
|
if gg.__len__() > 30:
|
||
|
|
gg = gg[:30]
|
||
|
|
grippers = gg.to_open3d_geometry_list()
|
||
|
|
cloud = o3d.geometry.PointCloud()
|
||
|
|
cloud.points = o3d.utility.Vector3dVector(pc.astype(np.float32))
|
||
|
|
o3d.visualization.draw_geometries([cloud, *grippers])
|
||
|
|
|
||
|
|
# # Example code for execution
|
||
|
|
# g = gg[0]
|
||
|
|
# translation = g.translation
|
||
|
|
# rotation = g.rotation_matrix
|
||
|
|
|
||
|
|
# pose = translation_rotation_2_matrix(translation,rotation) #transform into 4x4 matrix, should be easy
|
||
|
|
# # Transform the grasp pose from camera frame to robot coordinate, implement according to your robot configuration
|
||
|
|
# tcp_pose = Camera_To_Robot(pose)
|
||
|
|
|
||
|
|
|
||
|
|
# tcp_ready_pose = copy.deepcopy(tcp_pose)
|
||
|
|
# tcp_ready_pose[:3, 3] = tcp_ready_pose[:3, 3] - 0.1 * tcp_ready_pose[:3, 2] # The ready pose is backward along the actual grasp pose by 10cm to avoid collision
|
||
|
|
|
||
|
|
# tcp_away_pose = copy.deepcopy(tcp_pose)
|
||
|
|
|
||
|
|
# # to avoid the gripper rotate around the z_{tcp} axis in the clock-wise direction.
|
||
|
|
# tcp_away_pose[3,:3] = np.array([0,0,-1], dtype=np.float64)
|
||
|
|
|
||
|
|
# # to avoid the object collide with the scene.
|
||
|
|
# tcp_away_pose[2,3] += 0.1
|
||
|
|
|
||
|
|
# # We rely on python-urx to send the tcp pose the ur5 arm, the package is available at https://github.com/SintefManufacturing/python-urx
|
||
|
|
# urx.movels([tcp_ready_pose, tcp_pose], acc = acc, vel = vel, radius = 0.05)
|
||
|
|
|
||
|
|
# # CLOSE_GRIPPER(), implement according to your robot configuration
|
||
|
|
# urx.movels([tcp_away_pose, self.throw_pose()], acc = 1.2 * acc, vel = 1.2 * vel, radius = 0.05, wait=False)
|
||
|
|
|