From 8aac9a75f1622fc0d0ae5b26cd3b3bcc60d3ee84 Mon Sep 17 00:00:00 2001 From: cleong110 <122366389+cleong110@users.noreply.github.com> Date: Fri, 22 Nov 2024 17:56:44 -0500 Subject: [PATCH] Feature/recursive videos to poses (#126) * CDL: minor doc typo fix * videos_to_poses: add recursive and video_suffix options * CDL take out accidentally-added comment in docs * Check all vid extensions, account for name collisions, updated glob to check .pose files for speed * Take out the THIS IS A PROBLEM from v0.1.md * Fix https://github.com/sign-language-processing/pose/pull/126#discussion_r1849066680 * directory.py: make pose_files a set. * Fix ugly if statement * Remove unneeded else * count instead of unneeded list, also refactoring SUPPORTED_VIDEO_FORMATS properly * CDL a bit more reformatting * Revert "Take out the THIS IS A PROBLEM from v0.1.md" This reverts commit 87346c008c65e4a42d9e1d107110ad6c64b52fe9. * Update v0.1.md --------- Co-authored-by: Colin Leong <--unset> Co-authored-by: Amit Moryossef --- README.md | 1 + docs/specs/v0.1.md | 1 + src/python/pose_format/bin/directory.py | 169 ++++++++++++++++++++---- 3 files changed, 143 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 2fc94a6..020b8d2 100644 --- a/README.md +++ b/README.md @@ -222,6 +222,7 @@ Alternatively, use a different testing framework to run tests, such as pytest. T * Or employ pytest: ```bash +# From src/python directory pytest . # or for a single file pytest pose_format/tensorflow/masked/tensor_test.py diff --git a/docs/specs/v0.1.md b/docs/specs/v0.1.md index b3ea36b..2ee3aaa 100644 --- a/docs/specs/v0.1.md +++ b/docs/specs/v0.1.md @@ -24,6 +24,7 @@ \[`unsigned short` Green] \[`unsigned short` Blue] + # Body \[`unsined short` FPS] \[`unsined short` Number of frames] # THIS IS A PROBLEM diff --git a/src/python/pose_format/bin/directory.py b/src/python/pose_format/bin/directory.py index 8d92a74..60205f0 100644 --- a/src/python/pose_format/bin/directory.py +++ b/src/python/pose_format/bin/directory.py @@ -1,45 +1,158 @@ import argparse -import os - +from pathlib import Path from pose_format.bin.pose_estimation import pose_video, parse_additional_config +from typing import List +import logging from tqdm import tqdm +# Note: untested other than .mp4. Support for .webm may have issues: https://github.com/sign-language-processing/pose/pull/126 +SUPPORTED_VIDEO_FORMATS = [".mp4", ".mov", ".avi", ".mkv", ".flv", ".wmv", ".webm"] + + +def find_videos_with_missing_pose_files( + directory: Path, + video_suffixes: List[str] = None, + recursive: bool = False, + keep_video_suffixes: bool = False, +) -> List[Path]: + """ + Finds videos with missing .pose files. + + Parameters + ---------- + directory: Path, + Directory to search for videos in. + video_suffixes: List[str], optional + Suffixes to look for, e.g. [".mp4", ".webm"]. If None, will use _SUPPORTED_VIDEO_FORMATS + recursive: bool, optional + Whether to look for video files recursively, or just the top-level. Defaults to false. + keep_video_suffixes: bool, optional + If true, when checking will append .pose suffix (e.g. foo.mp4->foo.mp4.pose, foo.webm->foo.webm.pose), + If false, will replace it (foo.mp4 becomes foo.pose, and foo.webm ALSO becomes foo.pose). + Default is false, which can cause name collisions. + + Returns + ------- + List[Path] + List of video paths without corresponding .pose files. + """ + + # Prevents the common gotcha with mutable default arg lists: + # https://docs.python-guide.org/writing/gotchas/#mutable-default-arguments + if video_suffixes is None: + video_suffixes = SUPPORTED_VIDEO_FORMATS + + glob_method = getattr(directory, "rglob" if recursive else "glob") + all_files = list(glob_method(f"*")) + video_files = [path for path in all_files if path.suffix in video_suffixes] + pose_files = {path for path in all_files if path.suffix == ".pose"} -def removesuffix(text: str, suffix: str): - if text.endswith(suffix): - return text[:-len(suffix)] - else: - return text + videos_with_missing_pose_files = [] + for vid_path in video_files: + corresponding_pose = get_corresponding_pose_path(video_path=vid_path, keep_video_suffixes=keep_video_suffixes) + if corresponding_pose not in pose_files: + videos_with_missing_pose_files.append(vid_path) -def find_missing_pose_files(directory: str): - all_files = os.listdir(directory) - mp4_files = [f for f in all_files if f.endswith(".mp4")] - pose_files = {removesuffix(f, ".pose") for f in all_files if f.endswith(".pose")} - missing_pose_files = [] + return videos_with_missing_pose_files - for mp4_file in mp4_files: - base_name = removesuffix(mp4_file, ".mp4") - if base_name not in pose_files: - missing_pose_files.append(os.path.join(directory, mp4_file)) - return sorted(missing_pose_files) +def get_corresponding_pose_path(video_path: Path, keep_video_suffixes: bool = False) -> Path: + """ + Given a video path, and whether to keep the suffix, returns the expected corresponding path with .pose extension. + + Parameters + ---------- + video_path : Path + Path to a video file + keep_video_suffixes : bool, optional + Whether to keep suffix (e.g. foo.mp4 -> foo.mp4.pose) + or replace (foo.mp4->foo.pose). Defaults to replace. + + Returns + ------- + Path + pathlib Path + """ + if keep_video_suffixes: + return video_path.with_name(f"{video_path.name}.pose") + return video_path.with_suffix(".pose") def main(): parser = argparse.ArgumentParser() - parser.add_argument('--format', - choices=['mediapipe'], - default='mediapipe', - type=str, - help='type of pose estimation to use') - parser.add_argument("--directory", type=str, required=True) - parser.add_argument('--additional-config', type=str, help='additional configuration for the pose estimator') + parser.add_argument( + "-f", + "--format", + choices=["mediapipe"], + default="mediapipe", + type=str, + help="type of pose estimation to use", + ) + parser.add_argument( + "-d", + "--directory", + type=Path, + required=True, + help="Directory to search for videos in", + ) + parser.add_argument( + "-r", + "--recursive", + action="store_true", + help="Whether to search for videos recursively", + ) + parser.add_argument( + "--keep-video-suffixes", + action="store_true", + help="Whether to drop the video extension (output for foo.mp4 becomes foo.pose, and foo.webm ALSO becomes foo.pose) or append to it (foo.mp4 becomes foo.mp4.pose, foo.webm output is foo.webm.pose). If there are multiple videos with the same basename but different extensions, this will create a .pose file for each. Otherwise only the first video will be posed.", + ) + parser.add_argument( + "--video-suffixes", + type=str, + choices=SUPPORTED_VIDEO_FORMATS, + default=SUPPORTED_VIDEO_FORMATS, + help="Video extensions to search for. Defaults to searching for all supported.", + ) + parser.add_argument( + "--additional-config", + type=str, + help="additional configuration for the pose estimator", + ) args = parser.parse_args() - missing_pose_files = find_missing_pose_files(args.directory) + videos_with_missing_pose_files = find_videos_with_missing_pose_files( + args.directory, + video_suffixes=args.video_suffixes, + recursive=args.recursive, + keep_video_suffixes=args.keep_video_suffixes, + ) + + print(f"Found {len(videos_with_missing_pose_files)} videos missing pose files.") + + pose_files_that_will_be_created = {get_corresponding_pose_path(vid_path, args.keep_video_suffixes) for vid_path in videos_with_missing_pose_files} + + if len(pose_files_that_will_be_created) < len(videos_with_missing_pose_files): + continue_input = input( + f"With current naming strategy (without --keep-video-suffixes), name collisions will result in only {len(pose_files_that_will_be_created)} .pose files being created. Continue? [y/n]" + ) + if continue_input.lower() != "y": + print(f"Exiting. To keep video suffixes and avoid collisions, use --keep-video-suffixes") + exit() + additional_config = parse_additional_config(args.additional_config) - for mp4_path in tqdm(missing_pose_files): - pose_file_name = removesuffix(mp4_path, ".mp4") + ".pose" - pose_video(mp4_path, pose_file_name, args.format, additional_config) + pose_with_no_errors_count = 0 + + for vid_path in tqdm(videos_with_missing_pose_files): + try: + pose_path = get_corresponding_pose_path(video_path=vid_path, keep_video_suffixes=args.keep_video_suffixes) + if pose_path.is_file(): + print(f"Skipping {vid_path}, corresponding .pose file already created.") + continue + pose_video(vid_path, pose_path, args.format, additional_config) + pose_with_no_errors_count += 1 + except ValueError as e: + print(f"ValueError on {vid_path}") + logging.exception(e) + print(f"Successfully created pose files for {pose_with_no_errors_count}/{len(videos_with_missing_pose_files)} video files")