91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
|
|
|
|
def generate_merged_validation_json(args):
|
|
input_file = args.input_file
|
|
output_validation_file = args.output_validation_file
|
|
|
|
if args.output_train_file:
|
|
output_train_file = args.output_train_file
|
|
else:
|
|
base, ext = os.path.splitext(input_file)
|
|
output_train_file = f"{base}_train{ext}"
|
|
|
|
# read in input json
|
|
print(f"Reading from {input_file}")
|
|
with open(input_file, "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 input file: {num_elements}")
|
|
|
|
# randomly sample elements from the list
|
|
num_sample = min(args.num_elements, num_elements)
|
|
indices = set(random.sample(range(num_elements), num_sample))
|
|
|
|
sampled_elements = []
|
|
remaining_elements = []
|
|
|
|
for i in range(num_elements):
|
|
if i in indices:
|
|
sampled_elements.append(video2caption[i])
|
|
else:
|
|
remaining_elements.append(video2caption[i])
|
|
|
|
# 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": args.num_inference_steps,
|
|
"height": args.height,
|
|
"width": args.width,
|
|
"num_frames": args.num_frames
|
|
}
|
|
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_validation_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_validation_file}")
|
|
|
|
# Write the remaining JSON to the output train file
|
|
with open(output_train_file, "w") as f:
|
|
json.dump(remaining_elements, f, indent=2)
|
|
|
|
print(f"Saved remaining {len(remaining_elements)} entries to {output_train_file}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
# dataset_type: "merged"
|
|
parser.add_argument("--dataset_type", choices=["merged"], required=True)
|
|
parser.add_argument("--input_file", type=str, required=True, help="Path to input json file")
|
|
parser.add_argument("--output_validation_file", type=str, required=True, help="Path to output validation json file")
|
|
parser.add_argument("--output_train_file", type=str, help="Path to output train json file (remaining data). Defaults to {input_filename}_train.json")
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|