def __init__(self, regressor_checkpoint, smpl_dir, device = torch.device('cuda') , use_smplx = False): #For image transform transform_list = [ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] self.normalize_transform = transforms.Compose(transform_list) #Load Hand network self.opt = TestOptions().parse([]) #Default options self.opt.single_branch = True self.opt.main_encoder = "resnet50" # self.opt.data_root = "/home/hjoo/dropbox/hand_yu/data/" self.opt.model_root = "./extra_data" self.opt.smplx_model_file = os.path.join(smpl_dir,'SMPLX_NEUTRAL.pkl') self.opt.batchSize = 1 self.opt.phase = "test" self.opt.nThreads = 0 self.opt.which_epoch = -1 self.opt.checkpoint_path = regressor_checkpoint self.opt.serial_batches = True # no shuffle self.opt.no_flip = True # no flip self.opt.process_rank = -1 # self.opt.which_epoch = str(epoch) self.model_regressor = H3DWModel(self.opt) # if there is no specified checkpoint, then skip assert self.model_regressor.success_load, "Specificed checkpoints does not exists: {}".format(self.opt.checkpoint_path) self.model_regressor.eval()
def __init__(self, body_regressor_checkpoint, hand_regressor_checkpoint, smpl_dir, device=torch.device('cuda'), use_smplx=True): super().__init__('BodyMocap') self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') print("Loading Body Pose Estimator") self.__load_body_estimator() self.visualizer = Visualizer('opengl') self.frame_id = 0 #count frames parser = argparse.ArgumentParser() parser.add_argument("--rot90", default=False, type= bool, help="clockwise rotate 90 degrees") #parser.add_argument("--camera_topic", default="/logi_c922_2/image_rect_color", help="choose a topic as input image") parser.add_argument("--body_only", default=False, type= bool, help="detect only body and save its result") parser.add_argument("--result_path", default="/home/student/result/", help="choose a topic as input image") parser.add_argument("--save_result", default=False, help="save result or not") args = parser.parse_args() self.rot90 = args.rot90 #self.camera_topic = args.camera_topic self.body_only = args.body_only self.result_path = args.result_path self.save_result = args.save_result self.load = [0,0] self.angle_leg = 0 self.angle_trunk = 0 self.start = 0 self.angles = np.empty((1,20),dtype = float) self.body_side = np.empty((25,3),dtype = float) # Load parametric model (SMPLX or SMPL) if use_smplx: smplModelPath = smpl_dir + '/SMPLX_NEUTRAL.pkl' self.smpl = SMPLX(smpl_dir, batch_size=1, num_betas = 10, use_pca = False, create_transl=False).to(self.device) self.use_smplx = True else: smplModelPath = smpl_dir + '/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl' self.smpl = SMPL(smplModelPath, batch_size=1, create_transl=False).to(self.device) self.use_smplx = False #Load pre-trained neural network SMPL_MEAN_PARAMS = '/home/student/frankmocap/extra_data/body_module/data_from_spin/smpl_mean_params.npz' self.model_regressor = hmr(SMPL_MEAN_PARAMS).to(self.device) body_checkpoint = torch.load(body_regressor_checkpoint) self.model_regressor.load_state_dict(body_checkpoint['model'], strict=False) self.model_regressor.eval() #hand module init transform_list = [ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] self.normalize_transform = transforms.Compose(transform_list) #Load Hand network self.opt = TestOptions().parse([]) #Default options self.opt.single_branch = True self.opt.main_encoder = "resnet50" # self.opt.data_root = "/home/hjoo/dropbox/hand_yu/data/" self.opt.model_root = "/home/student/frankmocap/extra_data" self.opt.smplx_model_file = os.path.join(smpl_dir,'SMPLX_NEUTRAL.pkl') self.opt.batchSize = 1 self.opt.phase = "test" self.opt.nThreads = 0 self.opt.which_epoch = -1 self.opt.checkpoint_path = hand_regressor_checkpoint self.opt.serial_batches = True # no shuffle self.opt.no_flip = True # no flip self.opt.process_rank = -1 # self.opt.which_epoch = str(epoch) self.hand_model_regressor = H3DWModel(self.opt) # if there is no specified checkpoint, then skip assert self.hand_model_regressor.success_load, "Specificed checkpoints does not exists: {}".format(self.opt.checkpoint_path) self.hand_model_regressor.eval() self.hand_bbox_detector = HandBboxDetector('third_view', self.device) #subscriber and publisher initialization #input subscriber self.br = CvBridge() self.subscription_img = self.create_subscription(Image, '/side_img', self.callback_side,10) self.subscription_img = self.create_subscription(Image, '/front_img', self.callback_front,10) #output publisher self.publisher_pose = self.create_publisher(Image,'/pose',10) #images with keypoints annotation #self.publisher_keypoints = self.create_publisher(Float32MultiArray,'/keypoints',10) #keypoints coordinates self.publisher_risk = self.create_publisher(Int64,'/risk',10) #risk level self.publisher_angles = self.create_publisher(Float32MultiArray,'/angles',10)