(video_path, output_dir, fix_z)
| 141 | |
| 142 | |
| 143 | def get_pose3D(video_path, output_dir, fix_z): |
| 144 | args, _ = argparse.ArgumentParser().parse_known_args() |
| 145 | args.layers, args.channel, args.d_hid, args.token_dim, args.frames = 3, 512, 1024, 256, 243 |
| 146 | args.pad = (args.frames - 1) // 2 |
| 147 | args.previous_dir = 'checkpoint/pretrained/243_demo' |
| 148 | args.n_joints, args.out_joints = 17, 17 |
| 149 | |
| 150 | ## Reload |
| 151 | model = Model(args).cuda() |
| 152 | |
| 153 | model_dict = model.state_dict() |
| 154 | # Put the pretrained model in 'checkpoint/pretrained/243' |
| 155 | model_path = sorted(glob.glob(os.path.join(args.previous_dir, '*.pth')))[0] |
| 156 | |
| 157 | pre_dict = torch.load(model_path) |
| 158 | model_dict = model.state_dict() |
| 159 | state_dict = {k: v for k, v in pre_dict.items() if k in model_dict.keys()} |
| 160 | model_dict.update(state_dict) |
| 161 | model.load_state_dict(model_dict) |
| 162 | |
| 163 | model.eval() |
| 164 | |
| 165 | ## input |
| 166 | keypoints = np.load(output_dir + 'input_2D/input_keypoints_2d.npz', allow_pickle=True)['reconstruction'] |
| 167 | |
| 168 | cap = cv2.VideoCapture(video_path) |
| 169 | video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 170 | |
| 171 | ## 3D |
| 172 | print('\nGenerating 3D pose...') |
| 173 | output_3d_all = [] |
| 174 | for i in tqdm(range(video_length)): |
| 175 | ret, img = cap.read() |
| 176 | img_size = img.shape |
| 177 | |
| 178 | ## input frames |
| 179 | start = max(0, i - args.pad) |
| 180 | end = min(i + args.pad, len(keypoints[0])-1) |
| 181 | |
| 182 | input_2D_no = keypoints[0][start:end+1] |
| 183 | |
| 184 | left_pad, right_pad = 0, 0 |
| 185 | if input_2D_no.shape[0] != args.frames: |
| 186 | if i < args.pad: |
| 187 | left_pad = args.pad - i |
| 188 | if i > len(keypoints[0]) - args.pad - 1: |
| 189 | right_pad = i + args.pad - (len(keypoints[0]) - 1) |
| 190 | |
| 191 | input_2D_no = np.pad(input_2D_no, ((left_pad, right_pad), (0, 0), (0, 0)), 'edge') |
| 192 | |
| 193 | joints_left = [4, 5, 6, 11, 12, 13] |
| 194 | joints_right = [1, 2, 3, 14, 15, 16] |
| 195 | |
| 196 | input_2D = normalize_screen_coordinates(input_2D_no, w=img_size[1], h=img_size[0]) |
| 197 | |
| 198 | input_2D_aug = copy.deepcopy(input_2D) |
| 199 | input_2D_aug[ :, :, 0] *= -1 |
| 200 | input_2D_aug[ :, joints_left + joints_right] = input_2D_aug[ :, joints_right + joints_left] |
no test coverage detected