import os
import shutil
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--root_path', type=str, default='datasets/ADE20K_2021_17_01')
parser.add_argument('--save_path', type=str, default='datasets/ADE20K_val')
args = parser.parse_args()

if __name__ == "__main__":
    val_img_path = os.path.join(args.root_path, 'images/ADE/validation')
    filenames = os.listdir(val_img_path)
    os.makedirs(args.save_path, exist_ok=True)
    
    for filename in filenames:
        file_path = os.path.join(val_img_path, filename)
        if not os.path.isdir(file_path):
            continue
        child_files = os.listdir(file_path)
        for cfile in child_files:
            cfile_path = os.path.join(file_path, cfile)
            if not os.path.isdir(cfile_path):
                continue
            imgnames = os.listdir(cfile_path)
            for img_name in imgnames:
                if not img_name.endswith('.jpg'):
                    continue
                source_path = os.path.join(cfile_path, img_name)
                target_path = os.path.join(args.save_path, img_name)
                print("move {} to {}".format(source_path, target_path))
                shutil.move(source_path, target_path)
                
    print('### process end')