(args)
| 12 | from loss import joints_2d_loss, pose_temporal_loss, get_loss_weights |
| 13 | |
| 14 | def main(args): |
| 15 | device = torch.device("cuda:0") |
| 16 | seq = args.seq |
| 17 | gender = args.gender |
| 18 | DIR = './raw_data' |
| 19 | img_dir = f'{DIR}/{seq}/frames' |
| 20 | romp_file_dir = f'{DIR}/{seq}/ROMP' |
| 21 | img_paths = sorted(glob.glob(f"{img_dir}/*.png")) |
| 22 | romp_file_paths = sorted(glob.glob(f"{romp_file_dir}/*.npz")) |
| 23 | |
| 24 | from smplx import SMPL |
| 25 | smpl_model = SMPL('../code/lib/smpl/smpl_model', gender=gender).to(device) |
| 26 | |
| 27 | input_img = cv2.imread(img_paths[0]) |
| 28 | if args.source == 'custom': |
| 29 | focal_length = max(input_img.shape[0], input_img.shape[1]) |
| 30 | cam_intrinsics = np.array([[focal_length, 0., input_img.shape[1]//2], |
| 31 | [0., focal_length, input_img.shape[0]//2], |
| 32 | [0., 0., 1.]]) |
| 33 | elif args.source == 'neuman': |
| 34 | NeuMan_DIR = '' # path to NeuMan dataset |
| 35 | with open(f'{NeuMan_DIR}/{seq}/sparse/cameras.txt') as f: |
| 36 | lines = f.readlines() |
| 37 | cam_params = lines[3].split() |
| 38 | cam_intrinsics = np.array([[float(cam_params[4]), 0., float(cam_params[6])], |
| 39 | [0., float(cam_params[5]), float(cam_params[7])], |
| 40 | [0., 0., 1.]]) |
| 41 | elif args.source == 'deepcap': |
| 42 | DeepCap_DIR = '' # path to DeepCap dataset |
| 43 | with open(f'{DeepCap_DIR}/monocularCalibrationBM.calibration') as f: |
| 44 | lines = f.readlines() |
| 45 | |
| 46 | cam_params = lines[5].split() |
| 47 | cam_intrinsics = np.array([[float(cam_params[1]), 0., float(cam_params[3])], |
| 48 | [0., float(cam_params[6]), float(cam_params[7])], |
| 49 | [0., 0., 1.]]) |
| 50 | else: |
| 51 | print('Please specify the source of the dataset (custom, neuman, deepcap). We will continue to update the sources in the future.') |
| 52 | raise NotImplementedError |
| 53 | renderer = Renderer(img_size = [input_img.shape[0], input_img.shape[1]], cam_intrinsic=cam_intrinsics) |
| 54 | |
| 55 | if args.mode == 'mask': |
| 56 | if not os.path.exists(f'{DIR}/{seq}/init_mask'): |
| 57 | os.makedirs(f'{DIR}/{seq}/init_mask') |
| 58 | elif args.mode == 'refine': |
| 59 | if not os.path.exists(f'{DIR}/{seq}/init_refined_smpl'): |
| 60 | os.makedirs(f'{DIR}/{seq}/init_refined_smpl') |
| 61 | if not os.path.exists(f'{DIR}/{seq}/init_refined_mask'): |
| 62 | os.makedirs(f'{DIR}/{seq}/init_refined_mask') |
| 63 | if not os.path.exists(f'{DIR}/{seq}/init_refined_smpl_files'): |
| 64 | os.makedirs(f'{DIR}/{seq}/init_refined_smpl_files') |
| 65 | openpose_dir = f'{DIR}/{seq}/openpose' |
| 66 | openpose_paths = sorted(glob.glob(f"{openpose_dir}/*.npy")) |
| 67 | opt_num_iters=150 |
| 68 | weight_dict = get_loss_weights() |
| 69 | cam = PerspectiveCamera(focal_length_x=torch.tensor(cam_intrinsics[0, 0], dtype=torch.float32), |
| 70 | focal_length_y=torch.tensor(cam_intrinsics[1, 1], dtype=torch.float32), |
| 71 | center=torch.tensor(cam_intrinsics[0:2, 2]).unsqueeze(0)).to(device) |
no test coverage detected