File size: 1,237 Bytes
15e1e77 |
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 |
import argparse
import json
import os
import sys
def main(base_dir: str, output_file: str):
list_classes = ["A", "B", "C", "D"]
output_list = []
for _class in list_classes:
data_dir = os.path.join(base_dir, _class)
list_files = os.listdir(data_dir)
if _class == "A":
_label = [1, 0, 0, 0]
elif _class == "B":
_label = [0, 1, 0, 0]
elif _class == "C":
_label = [0, 0, 1, 0]
elif _class == "D":
_label = [0, 0, 0, 1]
for _file in list_files:
_out = {"image": os.path.join(data_dir, _file), "label": _label}
output_list.append(_out)
data_dict = {"Test": output_list}
fid = open(output_file, "w")
json.dump(data_dict, fid, indent=1)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="")
parser.add_argument("-base_dir", "--base_dir", default="sample_data", help="dir of dataset")
parser.add_argument(
"-output_file", "--output_file", default="configs/sample_image_data.json", help="output file name"
)
parser_args, _ = parser.parse_known_args(sys.argv)
main(base_dir=parser_args.base_dir, output_file=parser_args.output_file)
|