Helper function for cropping face and extracting landmarks.
(
org_path: Path,
mask_path: Path,
save_path: Path,
mode: str,
num_frames: int,
stride: int,
face_predictor: dlib.shape_predictor,
face_detector: dlib.fhog_object_detector,
margin: float = 0.5,
visualization: bool = False
)
| 247 | face_predictor = dlib.shape_predictor(predictor_path) |
| 248 | |
| 249 | def facecrop( |
| 250 | org_path: Path, |
| 251 | mask_path: Path, |
| 252 | save_path: Path, |
| 253 | mode: str, |
| 254 | num_frames: int, |
| 255 | stride: int, |
| 256 | face_predictor: dlib.shape_predictor, |
| 257 | face_detector: dlib.fhog_object_detector, |
| 258 | margin: float = 0.5, |
| 259 | visualization: bool = False |
| 260 | ) -> None: |
| 261 | """ |
| 262 | Helper function for cropping face and extracting landmarks. |
| 263 | """ |
| 264 | |
| 265 | # Open the video file |
| 266 | assert org_path.exists(), f"Video file {org_path} does not exist." |
| 267 | cap_org = cv2.VideoCapture(str(org_path)) |
| 268 | if not cap_org.isOpened(): |
| 269 | logger.error(f"Failed to open {org_path}") |
| 270 | return |
| 271 | |
| 272 | if mask_path is not None: |
| 273 | cap_mask = cv2.VideoCapture(str(mask_path)) |
| 274 | if not cap_mask.isOpened(): |
| 275 | logger.error(f"Failed to open {mask_path}") |
| 276 | return |
| 277 | |
| 278 | # Get the number of frames in the video |
| 279 | frame_count_org = int(cap_org.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 280 | |
| 281 | # Get the mode |
| 282 | if mode == 'fixed_num_frames': |
| 283 | # Get the frame rate of the video by dividing the number of frames by the duration (same interval between frames) |
| 284 | frame_idxs = np.linspace(0, frame_count_org - 1, num_frames, endpoint=True, dtype=int) |
| 285 | elif mode == 'fixed_stride': |
| 286 | # Get the frame rate of the video by dividing the number of frames by the duration (same interval between frames) |
| 287 | frame_idxs = np.arange(0, frame_count_org, stride, dtype=int) |
| 288 | |
| 289 | # Iterate through the frames |
| 290 | for cnt_frame in range(frame_count_org): |
| 291 | ret_org, frame_org = cap_org.read() |
| 292 | if mask_path is not None: |
| 293 | ret_mask, frame_mask = cap_mask.read() |
| 294 | else: |
| 295 | frame_mask = None |
| 296 | height, width = frame_org.shape[:-1] |
| 297 | |
| 298 | # Save original extracted frames |
| 299 | frame_path__ = save_path / 'frames_wocropface' / org_path.stem |
| 300 | frame_path__.mkdir(parents=True, exist_ok=True) |
| 301 | # Save |
| 302 | ori_frame_path = frame_path__ / f"{cnt_frame:03d}.png" |
| 303 | if not ori_frame_path.is_file(): |
| 304 | cv2.imwrite(str(ori_frame_path), frame_org) |
| 305 | |
| 306 | # Check if the frame was successfully read |
no test coverage detected