File size: 2,244 Bytes
ec0a9aa | 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 | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import glob
import json
import os
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset_path", type=str, required=True)
parser.add_argument("--save_path", type=str, required=True)
parser.add_argument("--output_path", type=str, required=True)
return parser.parse_args()
def main():
args = parse_args()
dataset_path = args.dataset_path
save_path = args.save_path
output_path = args.output_path
input_files = glob.glob(os.path.join(dataset_path, "*.jpg")) + glob.glob(os.path.join(dataset_path, "*.png"))
output_json = []
for input_file in input_files:
print(input_file)
prompt_file = input_file.replace(".jpg", ".txt").replace(".png", ".txt")
if not os.path.exists(prompt_file):
# The "dream_gen_benchmark/gr1_object/12_Use the right hand to pick up bok choy from top black shelf to paper bag.png" is missing a . before .png
prompt_file = input_file.replace(".jpg", "..txt").replace(".png", "..txt")
output_json.append(
{
"input_video": input_file,
"prompt": open(prompt_file).read(),
"output_video": os.path.join(
save_path, os.path.basename(input_file.replace(".jpg", ".mp4").replace(".png", ".mp4"))
),
}
)
print(f"Saved {len(output_json)} items to {output_path}")
with open(output_path, "w") as f:
json.dump(output_json, f, indent=4)
if __name__ == "__main__":
main()
|