Support ply and pcd formats in sai-cli process (#20)

* Add option to export optimized pointclouds in sai-cli tool

* Fix typo

* Dont create images folder if exporting ply or pcd pointcloud

* Cleanup code

* Revert unintentional changes

Author: kaatrasa

python/cli/process/process.py

@@ -1,15 +1,15 @@
 """
 Post-process data in Spectacular AI format and convert it to input
-for NeRF or Gaussian Splatting methods.
+for NeRF or Gaussian Splatting methods, or export optimized pointclouds in ply and pcd formats.
 """
 
 # --- The following mechanism allows using this both as a stand-alone
 # script and as a subcommand in sai-cli.
 
 def define_args(parser):
     parser.add_argument("input", help="Path to folder with session to process")
-    parser.add_argument("output", help="Output folder")
-    parser.add_argument('--format', choices=['taichi', 'nerfstudio'], default='nerfstudio', help='Output format')
+    parser.add_argument("output", help="Output folder, or filename with .ply or .pcd extension if exporting pointcloud")
+    parser.add_argument('--format', choices=['taichi', 'nerfstudio'], default='nerfstudio', help='Output format.')
     parser.add_argument("--cell_size", help="Dense point cloud decimation cell size (meters)", type=float, default=0.1)
     parser.add_argument("--distance_quantile", help="Max point distance filter quantile (0 = disabled)", type=float, default=0.99)
     parser.add_argument("--key_frame_distance", help="Minimum distance between keyframes (meters)", type=float, default=0.05)
@@ -33,10 +33,17 @@ def process(args):
     import json
     import os
     import shutil
+    import tempfile
     import numpy as np
     import pandas as pd
     from collections import OrderedDict
 
+    # Overwrite format if output is set to pointcloud
+    if args.output.endswith(".ply"):
+        args.format = "ply"
+    elif args.output.endswith(".pcd"):
+        args.format = "pcd"
+
     useMono = None
 
     def interpolate_missing_properties(df_source, df_query, k_nearest=3):
@@ -262,6 +269,10 @@ def onMappingOutput(output):
         if visualizer is not None:
             visualizer.onMappingOutput(output)
 
+        if args.format in ['ply', 'pcd']:
+            if output.finalMap: finalMapWritten = True
+            return
+
         if not output.finalMap:
             # New frames, let's save the images to disk
             for frameId in output.updatedKeyFrames:
@@ -288,7 +299,7 @@ def onMappingOutput(output):
                 img = undistortedFrame.image.toArray()
 
                 bgrImage = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-                fileName = f"{args.output}/tmp/frame_{frameId:05}.{args.image_format}"
+                fileName = f"{tmp_dir}/frame_{frameId:05}.{args.image_format}"
                 cv2.imwrite(fileName, bgrImage)
 
                 # Find colors for sparse features
@@ -307,7 +318,7 @@ def onMappingOutput(output):
                 if frameSet.depthFrame is not None and frameSet.depthFrame.image is not None and not useMono:
                     alignedDepth = frameSet.getAlignedDepthFrame(undistortedFrame)
                     depthData = alignedDepth.image.toArray()
-                    depthFrameName = f"{args.output}/tmp/depth_{frameId:05}.png"
+                    depthFrameName = f"{tmp_dir}/depth_{frameId:05}.png"
                     cv2.imwrite(depthFrameName, depthData)
 
                     DEPTH_PREVIEW = False
@@ -330,7 +341,7 @@ def onMappingOutput(output):
             sparsePointCloud = OrderedDict()
             imageSharpness = []
             for frameId in output.map.keyFrames:
-                imageSharpness.append((frameId, blurScore(f"{args.output}/tmp/frame_{frameId:05}.{args.image_format}")))
+                imageSharpness.append((frameId, blurScore(f"{tmp_dir}/frame_{frameId:05}.{args.image_format}")))
 
             # Look two images forward and two backwards, if current frame is blurriest, don't use it
             for i in range(len(imageSharpness)):
@@ -377,11 +388,11 @@ def onMappingOutput(output):
                     "camera_id": index # camera id, not used
                 }
 
-                oldImgName = f"{args.output}/tmp/frame_{frameId:05}.{args.image_format}"
+                oldImgName = f"{tmp_dir}/frame_{frameId:05}.{args.image_format}"
                 newImgName = f"{args.output}/images/frame_{index:05}.{args.image_format}"
                 os.rename(oldImgName, newImgName)
 
-                oldDepth = f"{args.output}/tmp/depth_{frameId:05}.png"
+                oldDepth = f"{tmp_dir}/depth_{frameId:05}.png"
                 newDepth = f"{args.output}/images/depth_{index:05}.png"
                 if os.path.exists(oldDepth):
                     os.rename(oldDepth, newDepth)
@@ -484,13 +495,6 @@ def detect_device_preset(input_dir):
                         if device: break
         return (device, cameras)
 
-    # Clear output dir
-    shutil.rmtree(f"{args.output}/images", ignore_errors=True)
-    os.makedirs(f"{args.output}/images", exist_ok=True)
-    tmp_dir = f"{args.output}/tmp"
-    tmp_input = f"{tmp_dir}/input"
-    copy_input_to_tmp_safe(args.input, tmp_input)
-
     config = {
         "maxMapSize": 0,
         "useSlam": True,
@@ -499,6 +503,17 @@ def detect_device_preset(input_dir):
         "icpVoxelSize": min(args.key_frame_distance, 0.1)
     }
 
+    if args.format in ['ply', 'pcd']:
+        config["mapSavePath"] = args.output
+    else:
+        # Clear output dir
+        shutil.rmtree(f"{args.output}/images", ignore_errors=True)
+        os.makedirs(f"{args.output}/images", exist_ok=True)
+
+    tmp_dir = tempfile.mkdtemp()
+    tmp_input = tempfile.mkdtemp()
+    copy_input_to_tmp_safe(args.input, tmp_input)
+
     device_preset, cameras = detect_device_preset(args.input)
 
     useMono = args.mono or (cameras != None and cameras == 1)
@@ -575,6 +590,11 @@ def detect_device_preset(input_dir):
     except:
         print(f"Failed to clean temporary directory, you can delete these files manually, they are no longer required: {tmp_dir}", flush=True)
 
+    try:
+        shutil.rmtree(tmp_input)
+    except:
+        print(f"Failed to clean temporary directory, you can delete these files manually, they are no longer required: {tmp_input}", flush=True)
+
     if not finalMapWritten:
         print('Mapping failed: no output generated')
         exit(1)
@@ -589,7 +609,7 @@ def detect_device_preset(input_dir):
         print(f"output-model-dir: data/{name}/output", flush=True)
         print(f"train-dataset-json-path: 'data/{name}/train.json'", flush=True)
         print(f"val-dataset-json-path: 'data/{name}/val.json'", flush=True)
-    elif args.format == 'nerfstudio':
+    else:
         print(f'output written to {args.output}', flush=True)
 
 if __name__ == '__main__':

python/cli/visualization/visualizer.py

@@ -315,6 +315,8 @@ def __initDisplay(self):
 
     def __close(self):
         assert(self.shouldQuit)
+        self.map.reset()
+        self.poseTrail.reset()
         if self.displayInitialized:
             self.displayInitialized = False
             pygame.quit()