63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
|
|
|
|
def generate_merged_validation_json(input_dir, output_file):
|
|
# read in video2caption.json
|
|
with open(os.path.join(input_dir, "video2caption_replace.json"), "r") as f:
|
|
video2caption = json.load(f)
|
|
|
|
# count how many elements are in the list
|
|
num_elements = len(video2caption)
|
|
print(f"Number of elements in video2caption.json: {num_elements}")
|
|
|
|
# randomly sample 64 elements from the list
|
|
sampled_elements = random.sample(video2caption, 64)
|
|
|
|
# Transform sampled elements into validation.json format
|
|
validation_data = []
|
|
for element in sampled_elements:
|
|
assert element.get("cap") is not None, f"Caption is None for element: {element}"
|
|
validation_entry = {
|
|
"caption": element["cap"],
|
|
"video_path": element.get("path", ""),
|
|
"num_inference_steps": 40,
|
|
"height": 480,
|
|
"width": 832,
|
|
"num_frames": 77
|
|
}
|
|
validation_data.append(validation_entry)
|
|
|
|
# Create the final validation structure
|
|
validation_json = {
|
|
"data": validation_data
|
|
}
|
|
|
|
# Write the validation JSON to the output file
|
|
with open(output_file, "w") as f:
|
|
json.dump(validation_json, f, indent=2)
|
|
|
|
print(f"Generated validation JSON with {len(validation_data)} entries and saved to {output_file}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
# dataset_type: "mixkit"
|
|
parser.add_argument("--dataset_type", choices=["merged"], required=True)
|
|
parser.add_argument("--input_dir", type=str, required=True)
|
|
parser.add_argument("--output_file", type=str, required=True)
|
|
parser.add_argument("--num_elements", type=int, default=64)
|
|
parser.add_argument("--num_frames", type=int, default=77)
|
|
parser.add_argument("--height", type=int, default=480)
|
|
parser.add_argument("--width", type=int, default=832)
|
|
parser.add_argument("--num_inference_steps", type=int, default=40)
|
|
args = parser.parse_args()
|
|
|
|
if args.dataset_type == "merged":
|
|
generate_merged_validation_json(args.input_dir, args.output_file)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |