| import argparse |
| import csv |
| import os |
|
|
| from tqdm import tqdm |
|
|
| |
| PREFIX = [ |
| "The video shows", |
| "The video captures", |
| "The video features", |
| "The video depicts", |
| "The video presents", |
| "The video features", |
| "The video is ", |
| "In the video,", |
| ] |
|
|
|
|
| def get_video_length(path): |
| import cv2 |
|
|
| cap = cv2.VideoCapture(path) |
| return int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
|
|
|
|
| def main(args): |
| input_path = args.input |
| output_path = args.output |
| if output_path is None: |
| name = os.path.basename(input_path) |
| name, ext = os.path.splitext(name) |
| if args.fmin is not None: |
| name += f"_fmin_{args.fmin}" |
| if args.fmax is not None: |
| name += f"_fmax_{args.fmax}" |
| if args.remove_empty_caption: |
| name += "_rec" |
| if args.remove_caption_prefix: |
| name += "_rcp" |
| if args.root is not None: |
| name += f"_root" |
| if args.relength: |
| name += "_relength" |
| output_path = os.path.join(os.path.dirname(input_path), name + ext) |
|
|
| with open(input_path, "r") as f: |
| reader = csv.reader(f) |
| data = list(reader) |
| print("Number of videos before filtering:", len(data)) |
|
|
| data_new = [] |
| for i, row in tqdm(enumerate(data)): |
| path = row[0] |
| caption = row[1] |
| n_frames = int(row[2]) |
| if args.fmin is not None and n_frames < args.fmin: |
| continue |
| if args.fmax is not None and n_frames > args.fmax: |
| continue |
| if args.remove_empty_caption and len(caption) == 0: |
| continue |
| if args.remove_caption_prefix: |
| for prefix in PREFIX: |
| if caption.startswith(prefix): |
| caption = caption[len(prefix) :].strip() |
| if caption[0].islower(): |
| caption = caption[0].upper() + caption[1:] |
| row[1] = caption |
| break |
| if args.root is not None: |
| row[0] = os.path.join(args.root, path) |
| if args.relength: |
| n_frames = get_video_length(row[0]) |
| row[2] = n_frames |
| data_new.append(row) |
|
|
| print("Number of videos after filtering:", len(data_new)) |
| with open(output_path, "w") as f: |
| writer = csv.writer(f) |
| writer.writerows(data_new) |
| print("Output saved to", output_path) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("input", type=str) |
| parser.add_argument("--output", type=str, default=None) |
| parser.add_argument("--fmin", type=int, default=None) |
| parser.add_argument("--fmax", type=int, default=None) |
| parser.add_argument("--root", type=str, default=None) |
| parser.add_argument("--remove-empty-caption", action="store_true") |
| parser.add_argument("--remove-caption-prefix", action="store_true") |
| parser.add_argument("--relength", action="store_true") |
| args = parser.parse_args() |
| main(args) |
|
|
