initial commit

datasets/build_skeleton_dataset.py

@@ -0,0 +1,39 @@
+#!/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()