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)

[((1499, 1524), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (1522, 1524), False, 'import argparse\n'), ((774, 804), 'os.path.join', 'os.path.join', (['root', 'label_path'], {}), '(root, label_path)\n', (786, 804), False, 'import os\n'), ((656, 686), 'os.path.join', 'os.path.join', (['root', 'label_path'], {}), '(root, label_path)\n', (668, 686), False, 'import os\n'), ((896, 933), 'os.path.join', 'os.path.join', (['root', 'image_path', 'image'], {}), '(root, image_path, image)\n', (908, 933), False, 'import os\n'), ((1310, 1341), 'os.path.join', 'os.path.join', (['root', 'output_file'], {}), '(root, output_file)\n', (1322, 1341), False, 'import os\n'), ((1420, 1451), 'os.path.join', 'os.path.join', (['root', 'output_file'], {}), '(root, output_file)\n', (1432, 1451), False, 'import os\n'), ((986, 1017), 'os.path.join', 'os.path.join', (['image_path', 'image'], {}), '(image_path, image)\n', (998, 1017), False, 'import os\n'), ((1019, 1050), 'os.path.join', 'os.path.join', (['label_path', 'label'], {}), '(label_path, label)\n', (1031, 1050), False, 'import os\n'), ((1124, 1161), 'os.path.join', 'os.path.join', (['root', 'image_path', 'image'], {}), '(root, image_path, image)\n', (1136, 1161), False, 'import os\n')]