Add python based batch data preparation.

data_generation/README.md

@@ -28,7 +28,7 @@ The code is only tested on Ubuntu 16.04 and 18.04.
 
 After compilation, binaries named `gen_3psdf_samples` and `batch_generate` will be generated.
 
-(1) `gen_3psdf_samples`: the binary generates 3PSDF sampling points using octree-based sampling. 
+### (1) `gen_3psdf_samples`: the binary generates 3PSDF sampling points using octree-based sampling. 
 
 Simply run `./gen_3psdf_samples` under the `build` folder will activate a demo setting that reconstructs `soldier_fight.obj` (a 3D character with both closed and open surfaces) under `/data` folder and saves (a) the generated 3PSDF samples for training (the`.sdf`file), (b) the mesh reconstructed from the computed 3PSDF field for debug purpose (the`.obj`file), and (c) the sampling point positions for debug purpose (the`.ply`file) to the `/output` folder.
 
@@ -55,7 +55,7 @@ To customize the usage of the program:
 `flag_writePLY`: flag of whether to generate `.ply` file that encodes the sampling points used in the `.sdf` training data. It is used for debug and visualization purpose. By default, it is set to 1 to activate the generation; it is disabled when set to 0. 
 
 
-(2) `batch_generate`: the binary that generates 3PSDF sampling points in batch. 
+### (2) `batch_generate`: the binary that generates 3PSDF sampling points in batch. 
 
 
 To customize the usage of the program:
@@ -82,6 +82,14 @@ To customize the usage of the program:
 `flag_writePLY`: flag of whether to generate `.ply` file that encodes the sampling points used in the `.sdf' training data. It is used for debug and visualization purpose. By default, it is set to 1 to activate the generation; it is disabled when set to 0. 
 
 
+### (3) `batch_generate.py`: a python 3.6+ wrapper for batch processing.
+
+Additionally keeps track of the exit status of the SDF generation in case it might fail for a mesh for some reason, in which case the binary `batch_generate` might exit and thus interrupt the generation process. Keeps track of the status in a `.json` and skips samples that have already been processed.
+
+Works with `objs` in subfolders and will preserve that structure for the output.
+
+Run `python batch_generate.py -h` for usage.
+
 ## Contacts
 
 If you have any questions, please contact Weikai Chen: <[email protected]>

data_generation/batch_generate.py

@@ -0,0 +1,63 @@
+import subprocess
+import os
+import json
+import argparse
+from pathlib import Path
+
+parser = argparse.ArgumentParser(description='Computes 3PSDF field from a given mesh and generates raw sampling data for network training.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+parser.add_argument("geo_dir", type=str, help='Folder containing input meshes that are ready to be batch processed.')
+parser.add_argument("sdf_dir", type=str, help='Folder for the generated `.sdf` files (raw data samples for 3PSDF training).')
+parser.add_argument("obj_dir", type=str, help='Folder for the reconstructed `.obj` files (reconstructed meshed from the output 3PSDF field).')
+parser.add_argument("ply_dir", type=str, help='Folder for the sampling points encoded in `.ply` format.')
+parser.add_argument("octree_depth", type=int, default=9, help='The depth of the octree that is used to generate sampling points. The larger the depth, the more accurate is the 3PSDF reconstruction - 6 (64^3), 7 (128^3), 8 (256^3) and 9 (512^3).')
+parser.add_argument('--writeSDF', action='store_true',  help='Flag of whether to generate `.sdf` file that encodes the raw data samples for 3PSDF learning.')
+parser.add_argument('--writeOBJ', action='store_true', help='Flag of whether to reconstruct the generated 3PSDF field to a mesh for debug purposes.')
+parser.add_argument('--writePLY', action='store_true', help='Flag of whether to generate `.ply` file that encodes the sampling points used in the `.sdf` training data.')
+parser.add_argument("--status_path", type=str, default="batch_generate_status.json", help='Filename for the status file to track the exit status of the data generation process. (1 is the successful exit status)')
+args = parser.parse_args()
+
+if os.path.isfile(args.status_path):
+    with open(args.status_path, "r") as f:
+        status_dict = json.load(f)
+else:
+    status_dict = {}
+
+obj_paths = list(Path(args.geo_dir).rglob('*.obj'))
+num_objs = len(obj_paths)
+for i, obj_path in enumerate(obj_paths):
+    fname = obj_path.stem
+    print(f"{i + 1}/{num_objs} {fname} ", flush=True, end="")
+    if fname in status_dict:
+        if status_dict[fname] == 1:
+            print("exists, skipping")
+            continue
+
+    subfolder = os.path.relpath(obj_path.parent, args.geo_dir)
+    sdf_dir = os.path.join(args.sdf_dir, subfolder)
+    if not os.path.isdir(sdf_dir) and args.writeSDF:
+        os.makedirs(sdf_dir)
+    obj_dir = os.path.join(args.obj_dir, subfolder)
+    if not os.path.isdir(obj_dir) and args.writeOBJ:
+        os.makedirs(obj_dir)
+    ply_dir = os.path.join(args.ply_dir, subfolder)
+    if not os.path.isdir(ply_dir) and args.writePLY:
+        os.makedirs(ply_dir)
+
+    status = subprocess.call(
+        [
+            "./build/gen_3psdf_samples",
+            str(obj_path),
+            os.path.join(sdf_dir, f"{fname}.sdf"),
+            os.path.join(obj_dir, f"{fname}.obj"),
+            os.path.join(ply_dir, f"{fname}.ply"),
+            str(args.octree_depth),
+            str(1) if args.writeSDF is True else str(0),
+            str(1) if args.writeOBJ is True else str(0),
+            str(1) if args.writePLY is True else str(0),
+        ],
+        stdout=subprocess.DEVNULL
+    )
+    print("OK" if status == 1 else "FAILED")
+    status_dict[fname] = status
+    with open(args.status_path, "w") as f:
+        json.dump(status_dict, f, indent=4)