def grab_new_objs(pkl_path, mano_path): rh_model = mano.load(model_path=mano_path, model_type='mano', num_pca_comps=45, batch_size=1, flat_hand_mean=True) rh_model_pkl = mano.load(model_path=mano_path, model_type='mano', num_pca_comps=15, batch_size=1, flat_hand_mean=False) all_samples = pickle.load(open(pkl_path, 'rb')) for idx, new_obj in enumerate(tqdm(all_samples)): print('idx', idx) ho = new_obj['ho_aug'] # obj_centroid = ho.obj_verts.mean(0) # ho.obj_verts = np.array(ho.obj_verts) - obj_centroid # ho.hand_verts = np.array(ho.hand_verts) - obj_centroid # ho.hand_mTc = np.array(ho.hand_mTc) # ho.hand_mTc[:3, 3] = ho.hand_mTc[:3, 3] - obj_centroid opt_dict = util.convert_pca15_aa45(ho, mano_model_in=rh_model_pkl, mano_model_out=rh_model)
def grab_new_objs(grabnet, objs_path, rot=True, n_samples=10, scale=1.): grabnet.coarse_net.eval() grabnet.refine_net.eval() rh_model = mano.load(model_path=grabnet.cfg.rhm_path, model_type='mano', num_pca_comps=45, batch_size=n_samples, flat_hand_mean=True).to(grabnet.device) grabnet.refine_net.rhm_train = rh_model grabnet.logger(f'################# \n' f'Grabbing the object!' ) bps = bps_torch(custom_basis=grabnet.bps) if not isinstance(objs_path, list): objs_path = [objs_path] for new_obj in objs_path: rand_rotdeg = np.random.random([n_samples, 3]) * np.array([360, 360, 360]) rand_rotmat = euler(rand_rotdeg) dorig = {'bps_object': [], 'verts_object': [], 'mesh_object': [], 'rotmat': []} for samples in range(n_samples): verts_obj, mesh_obj, rotmat = load_obj_verts(new_obj, rand_rotmat[samples], rndrotate=rot, scale=scale) bps_object = bps.encode(verts_obj, feature_type='dists')['dists'] dorig['bps_object'].append(bps_object.to(grabnet.device)) dorig['verts_object'].append(torch.from_numpy(verts_obj.astype(np.float32)).unsqueeze(0)) dorig['mesh_object'].append(mesh_obj) dorig['rotmat'].append(rotmat) obj_name = os.path.basename(new_obj) dorig['bps_object'] = torch.cat(dorig['bps_object']) dorig['verts_object'] = torch.cat(dorig['verts_object']) save_dir = os.path.join(grabnet.cfg.work_dir, 'grab_new_objects') grabnet.logger(f'#################\n' f' \n' f'Saving results for the {obj_name.upper()}' f' \n') gen_meshes = get_meshes(dorig=dorig, coarse_net=grabnet.coarse_net, refine_net=grabnet.refine_net, rh_model=rh_model, save=False, save_dir=save_dir ) torch.save(gen_meshes, 'data/grabnet_data/meshes.pt')
def __init__(self, data_dir, train=False, min_num_cont=1): start_time = time.time() self.dataset = pickle.load(open(data_dir, 'rb')) # Expensive step, can take up to 5 sec self.train = train self.aug_vert_jitter = 0.0005 # TODO, value? if 'num_verts_in_contact' in self.dataset[0]: print('Cutting samples less than {}. Was size {}'.format(min_num_cont, len(self.dataset))) self.dataset = [s for s in self.dataset if s['num_verts_in_contact'] >= min_num_cont] mano_path = '.' self.mano_model_out = mano.load(model_path=mano_path, model_type='mano', num_pca_comps=45, batch_size=1, flat_hand_mean=True) self.mano_model_in = mano.load(model_path=mano_path, model_type='mano', num_pca_comps=15, batch_size=1, flat_hand_mean=False) print('Dataset loaded in {:.2f} sec, {} samples'.format(time.time() - start_time, len(self.dataset)))
def __init__(self): self.obj_root = './models/HO3D_Object_models' obj_pc_dict, obj_face_dict, obj_scale_dict, obj_pc_resample_dict, obj_resample_faceid_dict = utils_HO3D_FPHA.load_objects_HO3D( self.obj_root) self.obj_pc_dict = obj_pc_dict self.obj_face_dict = obj_face_dict self.obj_scale_dict = obj_scale_dict self.obj_pc_resample_dict = obj_pc_resample_dict self.obj_resample_faceid_dict = obj_resample_faceid_dict self.nPoint = 3000 self.obj_list = list(self.obj_pc_dict.keys()) with torch.no_grad(): self.rh_mano = mano.load(model_path='./models/mano/MANO_RIGHT.pkl', model_type='mano', use_pca=False, num_pca_comps=45, batch_size=1, flat_hand_mean=True) self.hand_faces = self.rh_mano.faces.astype(np.int32).reshape( (-1, 3)) # [1538, 3], face triangle indexes
def __init__(self, cfg, inference=False, evaluate=False): self.dtype = torch.float32 torch.manual_seed(cfg.seed) starttime = datetime.now().replace(microsecond=0) makepath(cfg.work_dir, isfile=False) logger = makelogger( makepath(os.path.join(cfg.work_dir, '%s.log' % (cfg.expr_ID)), isfile=True)).info self.logger = logger summary_logdir = os.path.join(cfg.work_dir, 'summaries') self.swriter = SummaryWriter(log_dir=summary_logdir) logger('[%s] - Started training GrabNet, experiment code %s' % (cfg.expr_ID, starttime)) logger('tensorboard --logdir=%s' % summary_logdir) logger('Torch Version: %s\n' % torch.__version__) logger('Base dataset_dir is %s' % cfg.dataset_dir) # shutil.copy2(os.path.basename(sys.argv[0]), cfg.work_dir) use_cuda = torch.cuda.is_available() if use_cuda: torch.cuda.empty_cache() self.device = torch.device( "cuda:%d" % cfg.cuda_id if torch.cuda.is_available() else "cpu") gpu_brand = torch.cuda.get_device_name( cfg.cuda_id) if use_cuda else None gpu_count = torch.cuda.device_count() if cfg.use_multigpu else 1 if use_cuda: logger('Using %d CUDA cores [%s] for training!' % (gpu_count, gpu_brand)) self.data_info = {} self.load_data(cfg, inference) with torch.no_grad(): self.rhm_train = mano.load(model_path=cfg.rhm_path, model_type='mano', num_pca_comps=45, batch_size=cfg.batch_size // gpu_count, flat_hand_mean=True).to(self.device) self.coarse_net = CoarseNet().to(self.device) self.refine_net = RefineNet().to(self.device) self.LossL1 = torch.nn.L1Loss(reduction='mean') self.LossL2 = torch.nn.MSELoss(reduction='mean') if cfg.use_multigpu: self.coarse_net = nn.DataParallel(self.coarse_net) self.refine_net = nn.DataParallel(self.refine_net) logger("Training on Multiple GPU's") vars_cnet = [var[1] for var in self.coarse_net.named_parameters()] vars_rnet = [var[1] for var in self.refine_net.named_parameters()] cnet_n_params = sum(p.numel() for p in vars_cnet if p.requires_grad) rnet_n_params = sum(p.numel() for p in vars_rnet if p.requires_grad) logger('Total Trainable Parameters for CoarseNet is %2.2f M.' % ((cnet_n_params) * 1e-6)) logger('Total Trainable Parameters for RefineNet is %2.2f M.' % ((rnet_n_params) * 1e-6)) self.optimizer_cnet = optim.Adam(vars_cnet, lr=cfg.base_lr, weight_decay=cfg.reg_coef) self.optimizer_rnet = optim.Adam(vars_rnet, lr=cfg.base_lr, weight_decay=cfg.reg_coef) self.best_loss_cnet = np.inf self.best_loss_rnet = np.inf self.try_num = cfg.try_num self.epochs_completed = 0 self.cfg = cfg self.coarse_net.cfg = cfg if cfg.best_cnet is not None: self._get_cnet_model().load_state_dict(torch.load( cfg.best_cnet, map_location=self.device), strict=False) logger('Restored CoarseNet model from %s' % cfg.best_cnet) if cfg.best_rnet is not None: self._get_rnet_model().load_state_dict(torch.load( cfg.best_rnet, map_location=self.device), strict=False) logger('Restored RefineNet model from %s' % cfg.best_rnet) # weights for contact, penetration and distance losses self.vpe = torch.from_numpy(np.load(cfg.vpe_path)).to(self.device).to( torch.long) rh_f = torch.from_numpy(self.rhm_train.faces.astype(np.int32)).view( 1, -1, 3) self.rh_f = rh_f.repeat(self.cfg.batch_size, 1, 1).to(self.device).to(torch.long) v_weights = torch.from_numpy(np.load(cfg.c_weights_path)).to( torch.float32).to(self.device) v_weights2 = torch.pow(v_weights, 1.0 / 2.5) self.refine_net.v_weights = v_weights self.refine_net.v_weights2 = v_weights2 self.refine_net.rhm_train = self.rhm_train self.v_weights = v_weights self.v_weights2 = v_weights2 self.w_dist = torch.ones([self.cfg.batch_size, self.n_obj_verts]).to(self.device) self.contact_v = v_weights > 0.8
# load pre-trained model checkpoint_affordance = torch.load( args.affordance_model_path, map_location=torch.device('cpu'))['network'] affordance_model.load_state_dict(checkpoint_affordance) affordance_model = affordance_model.to(device) checkpoint_cmap = torch.load(args.cmap_model_path, map_location=torch.device('cpu'))['network'] cmap_model.load_state_dict(checkpoint_cmap) cmap_model = cmap_model.to(device) # dataset dataset = HO3D_diversity() dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=1) # mano hand model with torch.no_grad(): rh_mano = mano.load(model_path='./models/mano/MANO_RIGHT.pkl', model_type='mano', use_pca=True, num_pca_comps=45, batch_size=1, flat_hand_mean=True).to(device) rh_faces = torch.from_numpy(rh_mano.faces.astype(np.int32)).view( 1, -1, 3).to(device) # [1, 1538, 3], face indexes main(args, affordance_model, cmap_model, dataloader, device, rh_mano, rh_faces)
def grab_new_objs(grabnet, pkl_path, rot=True, n_samples=5, scale=1.): grabnet.coarse_net.eval() grabnet.refine_net.eval() rh_model = mano.load(model_path=grabnet.cfg.rhm_path, model_type='mano', num_pca_comps=45, batch_size=n_samples, flat_hand_mean=True).to(grabnet.device) rh_model_pkl = mano.load(model_path=grabnet.cfg.rhm_path, model_type='mano', num_pca_comps=15, batch_size=n_samples, flat_hand_mean=False).to(grabnet.device) grabnet.refine_net.rhm_train = rh_model grabnet.logger(f'################# \n' f'Colors Guide:' f' \n' f'Gray ---> GrabNet generated grasp\n') bps = bps_torch(custom_basis=grabnet.bps) all_samples = pickle.load(open(pkl_path, 'rb')) if args.vis: print('Shuffling!!!') random.shuffle(all_samples) all_samples = all_samples[:args.num] all_data = [] for idx, new_obj in enumerate(tqdm(all_samples)): print('idx', idx) ho = new_obj['ho_aug'] obj_centroid = ho.obj_verts.mean(0) ho.obj_verts = np.array(ho.obj_verts) - obj_centroid ho.hand_verts = np.array(ho.hand_verts) - obj_centroid ho.hand_mTc = np.array(ho.hand_mTc) ho.hand_mTc[:3, 3] = ho.hand_mTc[:3, 3] - obj_centroid rand_rotdeg = np.random.random([n_samples, 3]) * np.array([0, 0, 0]) rand_rotmat = euler(rand_rotdeg) dorig = { 'bps_object': [], 'verts_object': [], 'mesh_object': [], 'rotmat': [] } for samples in range(n_samples): verts_obj, mesh_obj, rotmat = load_obj_verts(ho, rand_rotmat[samples], rndrotate=rot, scale=scale) bps_object = bps.encode(verts_obj, feature_type='dists')['dists'] dorig['bps_object'].append(bps_object.to(grabnet.device)) dorig['verts_object'].append( torch.from_numpy(verts_obj.astype(np.float32)).unsqueeze(0)) dorig['mesh_object'].append(mesh_obj) dorig['rotmat'].append(rotmat) obj_name = 'test1' dorig['bps_object'] = torch.cat(dorig['bps_object']) dorig['verts_object'] = torch.cat(dorig['verts_object']) save_dir = os.path.join(grabnet.cfg.work_dir, 'grab_new_objects') # grabnet.logger(f'#################\n' # f' \n' # f'Showing results for the {obj_name.upper()}' # f' \n') verts_out, joints_out = vis_results(ho, dorig=dorig, coarse_net=grabnet.coarse_net, refine_net=grabnet.refine_net, rh_model=rh_model, save=False, save_dir=save_dir, rh_model_pkl=rh_model_pkl, vis=args.vis) ho.obj_verts = np.array(ho.obj_verts) + obj_centroid ho.hand_verts = np.array(ho.hand_verts) + obj_centroid ho.hand_mTc = np.array(ho.hand_mTc) ho.hand_mTc[:3, 3] = ho.hand_mTc[:3, 3] + obj_centroid verts_out = np.array( verts_out.detach().squeeze().numpy()) + obj_centroid joints_out = np.array( joints_out.detach().squeeze().numpy()) + obj_centroid new_ho = hand_object.HandObject() new_ho.load_from_verts(verts_out, new_obj['ho_gt'].obj_faces, new_obj['ho_gt'].obj_verts) all_data.append({ 'gt_ho': new_obj['ho_gt'], 'in_ho': new_obj['ho_aug'], 'out_verts': verts_out, 'out_joints': joints_out }) out_file = 'fitted_grabnet.pkl' print('Saving to {}. Len {}'.format(out_file, len(all_data))) pickle.dump(all_data, open(out_file, 'wb'))
def inference(grabnet): grabnet.coarse_net.eval() grabnet.refine_net.eval() ds_name = 'test' mesh_base = '/ps/scratch/grab/data/object_meshes/contact_meshes' ds_test = LoadData(dataset_dir=grabnet.cfg.dataset_dir, ds_name=ds_name) n_samples = 5 rh_model = mano.load(model_path=grabnet.cfg.rhm_path, model_type='mano', num_pca_comps=45, batch_size=n_samples, flat_hand_mean=True).to(grabnet.device) grabnet.refine_net.rhm_train = rh_model test_obj_names = np.unique(ds_test.frame_objs) grabnet.logger(f'################# \n' f'Colors Guide:' f' \n' f'Red ---> Reconstructed grasp - CoarseNet\n' f'Green ---> Reconstructed grasp - Refinent\n' f'Blue ---> Ground Truth Grasp\n' f'Pink ---> Generated grasp - CoarseNet\n' f'Gray ---> Generated grasp - RefineNet\n') for obj in test_obj_names: obj_frames = np.where(ds_test.frame_objs == obj)[0] rnd_frames = np.random.choice(obj_frames.shape[0], n_samples) obj_data = ds_test[obj_frames[rnd_frames]] frame_data = { k: obj_data[k].to(grabnet.device) for k in obj_data.keys() } obj_meshes = [] rotmats = [] for frame in range(n_samples): rot_mat = frame_data['root_orient_obj_rotmat'][frame].cpu().numpy( ).reshape(3, 3).T transl = frame_data['trans_obj'][frame].cpu().numpy() obj_mesh = Mesh(filename=os.path.join(mesh_base, obj + '.ply'), vc=name_to_rgb['yellow']) obj_mesh.rotate_vertices(rot_mat) obj_mesh.v += transl obj_meshes.append(obj_mesh) rotmats.append(rot_mat) frame_data['mesh_object'] = obj_meshes frame_data['rotmat'] = rotmats save_dir = os.path.join(grabnet.cfg.work_dir, 'test_grasp_results') grabnet.logger(f'#################\n' f' \n' f'Showing results for the {obj.upper()}' f' \n') vis_results(dorig=frame_data, coarse_net=grabnet.coarse_net, refine_net=grabnet.refine_net, rh_model=rh_model, show_rec=True, show_gen=True, save=False, save_dir=save_dir)