easyvolcap: adding more useful scripts

configs/specs/geometry.yaml

@@ -5,13 +5,14 @@ val_dataloader_cfg:
     dataset_cfg:
         type: GeometryDataset
         # skip_loading_images: True
+        cache_raw: False
         remove_outlier: False
-        vhull_thresh: 1.0
-        vhull_padding: 0.0 # padding could just be larger, this should be mostly fine
+        # vhull_thresh: 1.0
+        # vhull_padding: 0.0 # padding could just be larger, this should be mostly fine
 
-        vhull_count_factor: 1.0
-        vhull_thresh_factor: 1.0
-        vhull_ctof_factor: 20.0 # larger size?
+        # vhull_count_factor: 1.0
+        # vhull_thresh_factor: 1.0
+        # vhull_ctof_factor: 20.0 # larger size?
         view_sample: [0, null, 1] # use all views so that we can filter points to evaluate
 
 model_cfg:
@@ -24,7 +25,7 @@ runner_cfg:
     visualizer_cfg:
         type: GeometryVisualizer
         types: [MESH]
-        occ_thresh: 0.35 # leave out more points?
+        occ_thresh: 0.15 # leave out more points?
         result_dir: data/geometry
     evaluator_cfg:
         type: GeometryEvaluator

easyvolcap/dataloaders/datasets/geometry_dataset.py

@@ -28,8 +28,9 @@ def __init__(self,
 
         def carve_using_bytes(H, W, K, R, T, latent_index):
             if hasattr(self, 'mks_bytes'):
-                bytes = [self.mks_bytes[i * self.n_latents + latent_index] for i in range(len(H))]  # get mask bytes of this frame
-                msks = parallel_execution(bytes, normalize=True, action=load_image_from_bytes, sequential=True)
+                msks = [self.mks_bytes[i * self.n_latents + latent_index] for i in range(len(H))]  # get mask bytes of this frame
+                if not self.cache_raw:
+                    msks = parallel_execution(msks, normalize=True, action=load_image_from_bytes, sequential=True)
                 msks = to_tensor(msks)
 
                 # Fill blank canvas for each mask
@@ -78,6 +79,9 @@ def carve_using_bytes(H, W, K, R, T, latent_index):
             self.valid.append(valid)
             self.inds.append(inds)
 
+    def __len__(self):
+        return self.n_latents
+
     def __getitem__(self, index: int):
         output = self.get_metadata(index)
         output.xyz = self.xyz[output.latent_index]

easyvolcap/models/cameras/optimizable_camera.py

@@ -76,6 +76,15 @@ def __init__(self,
         if freeze_camera:
             freeze_module(self)
 
+        self.pre_handle = self._register_load_state_dict_pre_hook(self._load_state_dict_pre_hook)
+
+    def _load_state_dict_pre_hook(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+
+        # Historical reasons
+        if prefix + 'pose_resd' in state_dict:
+            if state_dict[prefix + 'pose_resd'].shape[0] == self.n_views and state_dict[prefix + 'pose_resd'].shape[1] == self.n_frames:
+                state_dict[prefix + 'pose_resd'] = state_dict[prefix + 'pose_resd'].transpose(0, 1)
+
     def forward_srcs(self, batch: dotdict):
         s_inds = batch.src_inds  # B, S, selected source views
         t_inds = batch.t_inds

easyvolcap/runners/visualizers/geometry_visualizer.py

@@ -36,47 +36,55 @@ def __init__(self,
                      Visualization.POINT.name: (lambda mesh, filename: export_pts(**mesh, filename=filename)),
                  },
                  verbose: bool = True,
-                 max_pending_pools: int = 100,  # maximum number of pending tasks in the thread pool
+                 pool_limit: int = 10,  # maximum number of pending tasks in the thread pool
 
                  occ_thresh: float = 0.5,
+                 sdf_thresh: float = 0.0,
                  **kwargs,
                  ):
         self.occ_thresh = occ_thresh
+        self.sdf_thresh = sdf_thresh
 
         result_dir = join(result_dir, cfg.exp_name)  # MARK: global configuration
         result_dir = join(result_dir, save_tag) if save_tag != '' else result_dir
         self.result_dir = result_dir
         self.types = [Visualization[t] for t in types]  # types of visualization
         self.exts = exts  # file extensions for each type of visualization
         self.exports = exports
+        self.verbose = verbose
 
         self.thread_pools: List[ThreadPool] = []
-        self.max_pending_pools = max_pending_pools
+        self.pool_limit = pool_limit
         self.geo_pattern = f'{{type}}/frame{{frame:04d}}_camera{{camera:04d}}{{ext}}'
 
         if verbose:
-            log(f'Visualization output: {yellow(join(result_dir, os.path.dirname(self.geo_pattern)))}')  # use yellow for output path
+            log(f'Visualization output: {yellow(join(result_dir, dirname(self.geo_pattern)))}')  # use yellow for output path
             log(f'Visualization types:', line(types))
 
     def generate_type(self, output: dotdict, batch: dotdict, type: Visualization = Visualization.MESH):
         if type == Visualization.MESH:
-            occ = output.occ
+            if 'sdf' in output:
+                occ = -output.sdf + 0.5
+                self.occ_thresh = -self.sdf_thresh + 0.5
+            else:
+                occ = output.occ
             voxel_size = batch.meta.voxel_size
             W, H, D = batch.meta.W[0].item(), batch.meta.H[0].item(), batch.meta.D[0].item()  # !: BATCH
-            cube = torch.zeros(np.prod(batch.valid.shape), dtype=occ.dtype, device='cpu')[None]  # 1, WHD
-            cube = multi_scatter_(cube[..., None], batch.inds.to('cpu', non_blocking=True), occ.to('cpu', non_blocking=True))  # dim = -2 # B, WHD, 1 assigned B, P, 1
+            cube = torch.full((np.prod(batch.valid.shape),), -10.0, dtype=occ.dtype, device='cpu')[None]  # 1, WHD
+            cube = multi_scatter_(cube[..., None], batch.inds.cpu(), occ.cpu())  # dim = -2 # B, WHD, 1 assigned B, P, 1
             cube = cube.view(-1, W, H, D)  # B, W, H, D
 
             # We leave the results on CPU but as tensors instead of numpy arrays
             torch.cuda.synchronize()  # some of the batched data are asynchronously moved to the cpu
-            verts, faces = mcubes.marching_cubes(cube.detach().cpu().float().numpy()[0], self.occ_thresh)
+            verts, faces = mcubes.marching_cubes(cube.float().numpy()[0], self.occ_thresh)
             verts = torch.as_tensor(verts, dtype=torch.float)[None]
             faces = torch.as_tensor(faces.astype(np.int32), dtype=torch.int)[None]
             verts = verts * voxel_size.to(verts.dtype) + batch.meta.bounds[:, 0].to(verts.dtype)  # !: BATCH
 
             mesh = dotdict()
             mesh.verts = verts
             mesh.faces = faces
+            log(f'Number of vertices: {verts.numel()}, faces: {faces.numel()} (frame: {batch.meta.frame_index.item()}, camera: {batch.meta.camera_index.item()})')
         else:
             raise NotImplementedError(f'Unimplemented visualization type: {type}')
         return mesh
@@ -119,11 +127,11 @@ def visualize(self, output: dotdict, batch: dotdict):
         return geo_stats
 
     def limit_thread_pools(self):
-        if len(self.thread_pools) > self.max_pending_pools:
-            for pool in self.thread_pools[:self.max_pending_pools]:
+        if len(self.thread_pools) > self.pool_limit:
+            for pool in self.thread_pools[:self.pool_limit]:
                 pool.close()
                 pool.join()
-            self.thread_pools = self.thread_pools[self.max_pending_pools:]
+            self.thread_pools = self.thread_pools[self.pool_limit:]
 
     def summarize(self):
         for pool in self.thread_pools:  # finish all pending taskes before generating videos

easyvolcap/utils/data_utils.py

@@ -105,6 +105,7 @@ def read_pfm(filename):
 
 def generate_video(result_str: str,
                    output: str,
+                   verbose: bool = False,
                    fps: int = 30,
                    crf: int = 17,
                    cqv: int = 19,
@@ -118,14 +119,19 @@ def generate_video(result_str: str,
     cmd = [
         'ffmpeg',
         '-hwaccel', hwaccel,
+    ] + ([
         '-hide_banner',
         '-loglevel', 'error',
+    ] if not verbose else []) + ([
         '-framerate', fps,
+    ] if fps > 0 else []) + ([
         '-f', 'image2',
         '-pattern_type', 'glob',
+    ] if not (splitext(result_str)[-1] or result_str.endswith('*')) else []) + ([
+        '-r', fps,
+    ] if fps > 0 else []) + [
         '-nostdin',  # otherwise you cannot chain commands together
         '-y',
-        '-r', fps,
         '-i', result_str,
         '-c:v', vcodec,
         '-preset', preset,

scripts/colmap/unflatten_dataset.py

@@ -1,3 +1,6 @@
+"""
+Find the images folder, rename it to images_flatten, and create a new images folder with the same structure as separated cameras
+"""
 # Rearrange images and camera parameters
 import shutil
 import argparse
@@ -6,9 +9,6 @@
 
 @catch_throw
 def main():
-    """
-    Find the images folder, rename it to images_flatten, and create a new images folder with the same structure as separated cameras
-    """
     parser = argparse.ArgumentParser()
     parser.add_argument('--data_root', default='data/zju/ipstage')
     parser.add_argument('--images_dirs', nargs='+', default=['images', 'masks'])

scripts/tools/0_to_00.py

@@ -6,21 +6,28 @@
 from easyvolcap.utils.easy_utils import read_camera, write_camera
 
 
+@catch_throw
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('input', default='data/mobile_stage/female')
+    parser.add_argument('--data_root', default='data/mobile_stage/female')
+    parser.add_argument('--dirs', default=['images', 'videos', 'masks', 'cameras'])
     args = parser.parse_args()
 
-    for cam in os.listdir(join(args.input, 'images')):
-        new_name = f"{int(cam):02d}"
-        if new_name != cam:
-            os.system(f'mv {join(args.input, "images", cam)} {join(args.input, "images", new_name)}')
+    for dir in args.dirs:
+        if not exists(join(args.data_root, dir)): continue
+        for cam in os.listdir(join(args.data_root, dir)):
+            idx, ext = splitext(cam)
+            idx = int(idx)
+            new_name = f"{idx:02d}{ext}"
+            if new_name != cam:
+                os.system(f'mv {join(args.data_root, dir, cam)} {join(args.data_root, dir, new_name)}')
 
-    cams = read_camera(join(args.input, 'intri.yml'), join(args.input, 'extri.yml'))
-    new_cams = {}
-    for cam in cams:
-        new_cams[f"{int(cam):02d}"] = cams[cam]
-    write_camera(new_cams, args.input)
+    if exists(join(args.data_root, 'intri.yml')) and exists(join(args.data_root, 'extri.yml')):
+        cams = read_camera(join(args.data_root, 'intri.yml'), join(args.data_root, 'extri.yml'))
+        new_cams = {}
+        for cam in cams:
+            new_cams[f"{int(cam):02d}"] = cams[cam]
+        write_camera(new_cams, args.data_root)
 
 
 if __name__ == '__main__':

scripts/tools/compress_videos.py

@@ -0,0 +1,27 @@
+"""
+Compress the video folder to something like libx265
+"""
+
+from easyvolcap.utils.console_utils import *
+from easyvolcap.utils.data_utils import generate_video
+
+
+@catch_throw
+def main():
+    args = dotdict(
+        data_root='data/bullet/final',
+        videos_dir='videos',
+        output_dir='videos_compressed',
+    )
+
+    args = dotdict(vars(build_parser(args, description=__doc__).parse_args()))
+    for video in sorted(os.listdir(join(args.data_root, args.videos_dir))):
+        if not video.endswith('.mp4'): continue
+        video_path = join(args.data_root, args.videos_dir, video)
+        output_path = join(args.data_root, args.output_dir, video)
+        os.makedirs(dirname(output_path), exist_ok=True)
+        generate_video(video_path, output_path, fps=-1, verbose=True)
+
+
+if __name__ == '__main__':
+    main()

scripts/tools/reorder_data.py

@@ -0,0 +1,26 @@
+import os
+import argparse
+from os.path import join
+
+from easyvolcap.utils.console_utils import *
+from easyvolcap.utils.easy_utils import read_camera, write_camera
+
+
+@catch_throw
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data_root', default='data/bullet/final')
+    parser.add_argument('--dirs', default=['images', 'videos', 'masks', 'cameras'])
+    args = parser.parse_args()
+
+    for dir in args.dirs:
+        if not exists(join(args.data_root, dir)): continue
+        for idx, cam in enumerate(sorted(os.listdir(join(args.data_root, dir)))):
+            _, ext = splitext(cam)
+            new_name = f"{idx:02d}{ext}"
+            if new_name != cam:
+                os.system(f'mv {join(args.data_root, dir, cam)} {join(args.data_root, dir, new_name)}')
+
+
+if __name__ == '__main__':
+    main()