Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,999 Bytes
e7d5680 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import argparse
import csv
import os
from tqdm import tqdm
# path, name, #frames
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)
|