File size: 3,151 Bytes
0dcccdd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import argparse
from pathlib import Path
import pandas as pd
from natsort import natsorted
from tqdm import tqdm
from .logger import logger
ALL_VIDEO_EXT = set(["mp4", "webm", "mkv", "avi", "flv", "mov"])
ALL_IMGAE_EXT = set(["png", "webp", "jpg", "jpeg", "bmp", "gif"])
def parse_args():
parser = argparse.ArgumentParser(description="Compute scores of uniform sampled frames from videos.")
parser.add_argument(
"--image_path_column",
type=str,
default="image_path",
help="The column contains the image path (an absolute path or a relative path w.r.t the image_folder).",
)
parser.add_argument("--image_folder", type=str, default=None, help="The video folder.")
parser.add_argument(
"--video_path_column",
type=str,
default="video_path",
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
)
parser.add_argument("--video_folder", type=str, default=None, help="The video folder.")
parser.add_argument("--saved_path", type=str, required=True, help="The save path to the output results (csv/jsonl).")
parser.add_argument("--recursive", action="store_true", help="Whether to search sub-folders recursively.")
args = parser.parse_args()
return args
def main():
args = parse_args()
if args.video_folder is None and args.image_folder is None:
raise ValueError("Either video_folder or image_folder should be specified in the arguments.")
if args.video_folder is not None and args.image_folder is not None:
raise ValueError("Both video_folder and image_folder can not be specified in the arguments at the same time.")
# Use the path name instead of the file name as video_path/image_path (unique ID).
if args.video_folder is not None:
video_path_list = []
video_folder = Path(args.video_folder)
for ext in tqdm(list(ALL_VIDEO_EXT)):
if args.recursive:
video_path_list += [str(file.relative_to(video_folder)) for file in video_folder.rglob(f"*.{ext}")]
else:
video_path_list += [str(file.relative_to(video_folder)) for file in video_folder.glob(f"*.{ext}")]
video_path_list = natsorted(video_path_list)
meta_file_df = pd.DataFrame({args.video_path_column: video_path_list})
if args.image_folder is not None:
image_path_list = []
image_folder = Path(args.image_folder)
for ext in tqdm(list(ALL_IMGAE_EXT)):
if args.recursive:
image_path_list += [str(file.relative_to(image_folder)) for file in image_folder.rglob(f"*.{ext}")]
else:
image_path_list += [str(file.relative_to(image_folder)) for file in image_folder.glob(f"*.{ext}")]
image_path_list = natsorted(image_path_list)
meta_file_df = pd.DataFrame({args.image_path_column: image_path_list})
logger.info(f"{len(meta_file_df)} files in total. Save the result to {args.saved_path}.")
meta_file_df.to_json(args.saved_path, orient="records", lines=True)
if __name__ == "__main__":
main() |