import argparse
import os

image_dir = 'image'
label_dir = 'label'
splits = ['train', 'val', 'test']
image_dirs = [
    'image/{}',
    'image/{}_crop'
]
label_dirs = [
    'label/{}/annotations',
    'label/{}/annotations_crop',
]


def generate(root):
    assert len(image_dirs) == len(label_dirs)

    for split in splits:
        for image_path, label_path in zip(image_dirs, label_dirs):
            image_path = image_path.format(split)
            label_path = label_path.format(split)

            if split != 'train' and image_path.endswith('_crop'):
                label_path = label_path.replace('_crop', '')

            if not os.path.exists(os.path.join(root, label_path)):
                continue

            lines = []
            for label in os.listdir(os.path.join(root, label_path)):
                image = label.replace('.png', '.jpg')

                if os.path.exists(os.path.join(root, image_path, image)):
                    lines.append('{} {}\n'.format(os.path.join(image_path, image), os.path.join(label_path, label)))
                else:
                    print('not found: {}'.format(os.path.join(root, image_path, image)))

            print(image_path, label_path, len(lines))

            output_file = '{}.lst'.format(image_path.split('/')[1])
            with open(os.path.join(root, output_file), 'w') as f:
                f.writelines(lines)

            print(f'Save to {os.path.join(root, output_file)}\n')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('root', type=str, help='path of dataset root')
    args = parser.parse_args()

    generate(args.root)

[((50, 13, 50, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((32, 36, 32, 66), 'os.path.join', 'os.path.join', ({(32, 49, 32, 53): 'root', (32, 55, 32, 65): 'label_path'}, {}), '(root, label_path)', False, 'import os\n'), ((28, 34, 28, 64), 'os.path.join', 'os.path.join', ({(28, 47, 28, 51): 'root', (28, 53, 28, 63): 'label_path'}, {}), '(root, label_path)', False, 'import os\n'), ((35, 34, 35, 71), 'os.path.join', 'os.path.join', ({(35, 47, 35, 51): 'root', (35, 53, 35, 63): 'image_path', (35, 65, 35, 70): 'image'}, {}), '(root, image_path, image)', False, 'import os\n'), ((43, 22, 43, 53), 'os.path.join', 'os.path.join', ({(43, 35, 43, 39): 'root', (43, 41, 43, 52): 'output_file'}, {}), '(root, output_file)', False, 'import os\n'), ((46, 29, 46, 60), 'os.path.join', 'os.path.join', ({(46, 42, 46, 46): 'root', (46, 48, 46, 59): 'output_file'}, {}), '(root, output_file)', False, 'import os\n'), ((36, 50, 36, 81), 'os.path.join', 'os.path.join', ({(36, 63, 36, 73): 'image_path', (36, 75, 36, 80): 'image'}, {}), '(image_path, image)', False, 'import os\n'), ((36, 83, 36, 114), 'os.path.join', 'os.path.join', ({(36, 96, 36, 106): 'label_path', (36, 108, 36, 113): 'label'}, {}), '(label_path, label)', False, 'import os\n'), ((38, 49, 38, 86), 'os.path.join', 'os.path.join', ({(38, 62, 38, 66): 'root', (38, 68, 38, 78): 'image_path', (38, 80, 38, 85): 'image'}, {}), '(root, image_path, image)', False, 'import os\n')]