#!/usr/bin/env python3
"""
Build dataset: given input images folder, extract skeletons and vectorize, save pairs.
"""
import argparse
from pathlib import Path
import subprocess


def main():
    p = argparse.ArgumentParser()
    p.add_argument('img_folder')
    p.add_argument('out_folder')
    p.add_argument('--invert', action='store_true')
    args = p.parse_args()

    out = Path(args.out_folder)
    out.mkdir(parents=True, exist_ok=True)

    imgs = list(Path(args.img_folder).glob('*.png')) + list(Path(args.img_folder).glob('*.jpg'))
    for im in imgs:
        od = out / im.stem
        od.mkdir(exist_ok=True)
        # call skeleton_extract
        cmd = ['python3', str(Path(__file__).parents[1] / 'scripts' / 'skeleton_extract.py'), str(im), str(od)]
        if args.invert:
            cmd.append('--invert')
        subprocess.run(cmd, check=True)
        # vectorize
        sk_png = od / (im.stem + '_sk.png')
        outjson = od / (im.stem + '_sk.json')
        cmd2 = ['python3', str(Path(__file__).parents[1] / 'scripts' / 'vectorize_skeleton.py'), str(sk_png), str(outjson)]
        subprocess.run(cmd2, check=True)

    print('Built dataset in', out)


if __name__ == '__main__':
    main()