import numpy as np
import os
import json
import argparse
dataset_name = "BirdCounting1"
root_path ="../data_MPKFNet/data/BirdCounting"


parser = argparse.ArgumentParser(description="DensityVideodemo.")
parser.add_argument("dataset_name", default='', type=str,
                    help='Video name.')

parser.add_argument(
    "opts",
    help="Modify config options using the command-line",
    default=None,
    nargs=argparse.REMAINDER,
)
args = parser.parse_args()
print(args)
dataset_name = args.dataset_name
datasetjsonpath = os.path.join(root_path ,dataset_name)
train = os.path.join(root_path,dataset_name,'train_data','images')
test = os.path.join(root_path,dataset_name,'test_data','images')
trainfile = os.listdir(train)
train_json = [os.path.join(train,file) for file in trainfile]
with open(datasetjsonpath + "/train_data.json",'w') as f:
    json.dump(train_json,f)

testfile = os.listdir(test)
test_json = [os.path.join(test,file) for file in testfile]
with open(datasetjsonpath + "/test_data.json",'w') as f:
    json.dump(test_json,f)

print("data is ok")