40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
#!/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()
|