adding control for metric computation selection

easyvolcap/runners/evaluators/volumetric_video_evaluator.py

@@ -1,14 +1,12 @@
-import torch
 import numpy as np
-from torch import nn
-import ujson as json
+from typing import List
 
 from easyvolcap.engine import EVALUATORS, cfg
 from easyvolcap.utils.console_utils import *
 from easyvolcap.utils.base_utils import dotdict
 from easyvolcap.utils.json_utils import serialize
 from easyvolcap.utils.data_utils import Visualization
-from easyvolcap.utils.metric_utils import psnr, ssim, lpips
+from easyvolcap.utils.metric_utils import psnr, ssim, lpips, Metrics
 from easyvolcap.runners.visualizers.volumetric_video_visualizer import VolumetricVideoVisualizer
 
 
@@ -19,13 +17,14 @@ def __init__(self,
                  result_dir: str = cfg.runner_cfg.visualizer_cfg.result_dir,  # MARK: GLOBAL
                  save_tag: str = cfg.runner_cfg.visualizer_cfg.save_tag,  # MARK: GLOBAL
                  metrics_file: str = 'metrics.json',
+                 compute_metrics: List[str] = [Metrics.PSNR.name, Metrics.SSIM.name, Metrics.LPIPS.name],
                  **kwargs,
                  ) -> None:
         super().__init__(verbose=False, result_dir=result_dir, save_tag=save_tag, **kwargs)
         self.skip_time_in_summary = skip_time_in_summary
         self.metrics = []
         self.metrics_file = metrics_file
-        self.compute_metrics = [psnr, ssim, lpips]
+        self.compute_metrics = [Metrics[m] for m in compute_metrics]
 
     def evaluate(self, output: dotdict, batch: dotdict):
         # TODO: This is a bit wasteful since the images are already generated by the visualizer

easyvolcap/runners/optimizers.py

@@ -80,7 +80,7 @@ def ConfigurableOptimizer(named_params: Iterator[Tuple[str, nn.Parameter]],
             dotdict(
                 params=[value],
                 lr=v_lr,
-                v_eps=v_eps,
+                eps=v_eps,
                 weight_decay=v_weight_decay,
                 name=item
             )

easyvolcap/runners/schedulers.py

@@ -8,6 +8,7 @@
 @SCHEDULERS.register_module()
 class MultiLR(_LRScheduler):
     def __init__(self, optimizer, decay_iter, scheduler_cfgs, last_epoch=-1, verbose=False):
+        raise ValueError('MultiLR has bugs! Please use other schedulers instead.')
         self.schedulers = dotdict()
         self.names = [param_group['name'] for param_group in optimizer.param_groups]
         # values = self._get_optimizer_lr(optimizer)

easyvolcap/utils/loss_utils.py

@@ -557,6 +557,4 @@ def compute_time_planes_smooth(embedding):
 
 def compute_ssim(x: torch.Tensor, y: torch.Tensor):
     from pytorch_msssim import ssim
-    x = x.permute(0, 3, 1, 2)
-    y = y.permute(0, 3, 1, 2)
     return ssim(x, y, data_range=1.0, win_size=11, win_sigma=1.5, K=(0.01, 0.03))

scripts/colmap/easymocap_to_colmap.py

@@ -1,70 +1,102 @@
 import os
-import argparse
-from os.path import join
-from easymocap.mytools.camera_utils import read_camera
+import cv2
 
-# fmt: off
-import sys
-sys.path.append('.')
+from easyvolcap.utils.console_utils import *
+from easymocap.mytools.camera_utils import read_camera
 from easyvolcap.utils.colmap_utils import write_cameras_binary, write_images_binary, Camera, Image, rotmat2qvec, qvec2rotmat, write_cameras_text, write_images_text
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--data_root',        default='data/my_zjumocap/my_313/bkgd')
-parser.add_argument('--output_dir',       default='text')
-parser.add_argument('-i', '--intri',      default='intri.yml')
-parser.add_argument('-e', '--extri',      default='extri.yml')
-parser.add_argument('--height',           default=1024)
-parser.add_argument('--width',            default=1024)
-args = parser.parse_args()
-args.output_dir = join(args.data_root, args.output_dir)
-args.intri = join(args.data_root, args.intri)
-args.extri = join(args.data_root, args.extri)
-# fmt: on
 
-cams = read_camera(args.intri, args.extri)
-assert 'basenames' in cams and len(cams['basenames'])
-basenames = sorted(cams['basenames'])
-os.makedirs(args.output_dir, exist_ok=True)
+@catch_throw
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data_root', default='data/my_zjumocap/my_313/bkgd')
+    parser.add_argument('--output_dir', default='colmap')
+    parser.add_argument('--image_dir', default='images_calib')
+    parser.add_argument('--image_ext', default='.jpg')
+    parser.add_argument('--mask_dir', default='masks')
+    parser.add_argument('--mask_ext', default='.jpg')
+    parser.add_argument('-i', '--intri', default='intri.yml')
+    parser.add_argument('-e', '--extri', default='extri.yml')
+    args = parser.parse_args()
+    args.output_dir = join(args.data_root, args.output_dir)
+    args.intri = join(args.data_root, args.intri)
+    args.extri = join(args.data_root, args.extri)
+
+    cams = read_camera(args.intri, args.extri)
+    assert 'basenames' in cams and len(cams['basenames'])
+    basenames = sorted(cams['basenames'])
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    cameras = {}
+    images = {}
+    sizes = {}
+    frames = os.listdir(join(args.data_root, args.image_dir, basenames[0]))
+    frames = sorted([x.split('.')[0] for x in frames])
+    for frame in frames:
+        output_dir = join(args.output_dir, frame)
+        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(join(output_dir, 'images'), exist_ok=True)
+        os.makedirs(join(output_dir, 'masks'), exist_ok=True)
+        os.makedirs(join(output_dir, 'sparse'), exist_ok=True)
+        for cam_id, cam_name in enumerate(basenames):
+            cam_id = cam_id + 1
+            # write images
+            # if os.path.exists(join(output_dir, 'images', f'{cam_id:02d}{args.image_ext}')):
+            os.remove(join(output_dir, 'images', f'{cam_id:02d}{args.image_ext}'))
+            # if os.path.exists(join(output_dir, 'masks', f'{cam_id:02d}{args.mask_ext}')):
+            os.remove(join(output_dir, 'masks', f'{cam_id:02d}{args.mask_ext}'))
+            os.symlink(join('../../..', args.image_dir, cam_name, f'{frame}{args.image_ext}'), join(output_dir, 'images', f'{cam_id:02d}{args.image_ext}'))
+            os.symlink(join('../../..', args.mask_dir, cam_name, f'{frame}{args.mask_ext}'), join(output_dir, 'masks', f'{cam_id:02d}{args.mask_ext}'))
+            # read image
+            if cam_name not in sizes.keys():
+                img = cv2.imread(join(output_dir, 'images', f'{cam_id:02d}{args.image_ext}'))
+                sizes[cam_name] = img.shape[:2]
+            # read camera
+            print(f'reading image and camera from: {cam_name}')
+            cam_dict = cams[cam_name]
+            K = cam_dict['K']
+            R = cam_dict['R']
+            T = cam_dict['T']
+            D = cam_dict['dist']  # !: losing k3 parameter
+            if 'H' in cam_dict.keys() or 'W' in cam_dict.keys():
+                cam_dict['H'] = sizes[cam_name][0]
+                cam_dict['W'] = sizes[cam_name][1]
+            if D.shape[0] == 1:
+                fx, fy, cx, cy, k1, k2, p1, p2, k3 = K[0, 0], K[1, 1], K[0, 2], K[1, 2], D[0, 0], D[0, 1], D[0, 2], D[0, 3], D[0, 4]
+            else:
+                fx, fy, cx, cy, k1, k2, p1, p2, k3 = K[0, 0], K[1, 1], K[0, 2], K[1, 2], D[0, 0], D[1, 0], D[2, 0], D[3, 0], D[4, 0]
+
+            params = [fx, fy, cx, cy, k1, k2, p1, p2]
+            camera = Camera(
+                id=cam_id,
+                model='OPENCV',
+                width=cam_dict['W'],
+                height=cam_dict['H'],
+                params=params
+            )
 
-cameras = {}
-images = {}
-for cam_id, cam_name in enumerate(basenames):
-    print(f'reading image and camera from: {cam_name}')
-    cam_dict = cams[cam_name]
-    K = cam_dict['K']
-    R = cam_dict['R']
-    T = cam_dict['T']
-    D = cam_dict['dist']  # !: losing k3 parameter
-    if D.shape[0] == 1:
-        fx, fy, cx, cy, k1, k2, p1, p2, k3 = K[0, 0], K[1, 1], K[0, 2], K[1, 2], D[0, 0], D[0, 1], D[0, 2], D[0, 3], D[0, 4]
-    else:
-        fx, fy, cx, cy, k1, k2, p1, p2, k3 = K[0, 0], K[1, 1], K[0, 2], K[1, 2], D[0, 0], D[1, 0], D[2, 0], D[3, 0], D[4, 0]
+            qvec = rotmat2qvec(R)
+            tvec = T.T[0]
+            name = f"{cam_id:02d}.jpg"
 
-    params = [fx, fy, cx, cy, k1, k2, p1, p2]
-    camera = Camera(
-        id=cam_id,
-        model='OPENCV',
-        width=args.width,
-        height=args.height,
-        params=params
-    )
+            image = Image(
+                id=cam_id,
+                qvec=qvec,
+                tvec=tvec,
+                camera_id=cam_id,
+                name=name,
+                xys=[],
+                point3D_ids=[]
+            )
 
-    qvec = rotmat2qvec(R)
-    tvec = T.T[0]
-    name = f"{cam_id:02d}.jpg"
+            cameras[cam_id] = camera
+            images[cam_id] = image
 
-    image = Image(
-        id=cam_id,
-        qvec=qvec,
-        tvec=tvec,
-        camera_id=cam_id,
-        name=name,
-        xys=[],
-        point3D_ids=[]
-    )
+        write_cameras_text(cameras, join(output_dir, 'sparse', 'cameras.txt'))
+        write_images_text(images, join(output_dir, 'sparse', 'images.txt'))
+        with open(join(output_dir, 'sparse', 'points3D.txt'), 'w') as f:
+            f.writelines(['# 3D point list with one line of data per point:\n'])
 
-    cameras[cam_id] = camera
-    images[cam_id] = image
 
-write_cameras_text(cameras, join(args.output_dir, 'cameras.txt'))
-write_images_text(images, join(args.output_dir, 'images.txt'))
+if __name__ == '__main__':
+    main()

scripts/preprocess/neural3dv_to_easyvolcap.py

@@ -4,15 +4,25 @@
 # Need a way to store the bound info in the camera parameters file
 
 import argparse
+import subprocess
 import numpy as np
-from os.path import join
+from glob import glob
+
 from easyvolcap.utils.console_utils import *
 from easyvolcap.utils.easy_utils import write_camera
 from easyvolcap.utils.math_utils import affine_inverse
 from easyvolcap.utils.data_utils import as_numpy_func, export_camera
 
 
 def main():
+    """
+    Assume ./data/neural3dv/XXX
+    in easyvolcap:
+    python3 scripts/preprocess/neural3dv_to_easyvolcap.py --only XXX
+    python3 scripts/colmap/easymocap_to_colmap.py --data_root data/neural3dv/XXX --image_dir images --output_dir colmap
+    in spg_colmap:
+    python3 sfm_renbody.py --root_dir ./data/neural3dv/XXX/colmap --colmap_path $PATHTOCOLMAP
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument('--neural3dv_root', type=str, default='data/neural3dv')
     parser.add_argument('--easyvolcap_root', type=str, default='data/neural3dv')
@@ -22,7 +32,24 @@ def main():
 
     scenes = os.listdir(args.neural3dv_root)
     scenes = [s for s in scenes if s in args.only]
-    for scene in tqdm(scenes):
+    # for scene in tqdm(scenes):
+    for scene in scenes:
+        videos = sorted(glob(join(args.neural3dv_root, scene, '*.mp4')))
+        for v in videos:
+            dirname = basename(v).split('.')[0][-2:]
+            if not exists(join(args.easyvolcap_root, scene, 'images', dirname)):
+                os.makedirs(join(args.easyvolcap_root, scene, 'images', dirname), exist_ok=True)
+                cmd = [
+                    'ffmpeg',
+                    '-i', v,
+                    '-vf', 'fps=30',
+                    '-q:v', '1',
+                    '-qmin', '1',
+                    '-start_number', '0',
+                    join(args.easyvolcap_root, scene, 'images', dirname) + '/%06d.jpg'
+                ]
+                subprocess.run(cmd, check=True)
+
         # https://github.com/kwea123/nerf_pl/blob/52aeb387da64a9ad9a0f914ea9b049ffc598b20c/datasets/llff.py#L177
         raw = np.load(join(args.neural3dv_root, scene, args.camera_pose), allow_pickle=True)  # 21, 17
         poses = raw[:, :15].reshape(-1, 3, 5)  # N, 3, 5

scripts/tools/copy_list_of_files.py

@@ -0,0 +1,35 @@
+"""
+Copy as list of files (folders), preserving directory structure
+"""
+from easyvolcap.utils.console_utils import *
+
+
+@catch_throw
+def main():
+    args = dotdict(
+        files=[
+            'easyvolcap/runners/evaluators/volumetric_video_evaluator.py',
+            'easyvolcap/runners/optimizers.py',
+            'easyvolcap/runners/schedulers.py',
+            'easyvolcap/utils/loss_utils.py',
+            'scripts/colmap/easymocap_to_colmap.py',
+            'scripts/preprocess/neural3dv_to_easyvolcap.py',
+            'scripts/tools/copy_list_of_files.py',
+        ],
+        source='.',
+        target='../easyvolcap-zju3dv',
+    )
+    args = dotdict(vars(build_parser(args, description=__doc__).parse_args()))
+
+    for f in tqdm(args.files):
+        src = join(args.source, f)
+        tar = join(args.target, f)
+        os.makedirs(dirname(tar), exist_ok=True)
+        if isfile(src):
+            shutil.copy(src, tar)
+        else:
+            shutil.copytree(src, tar)
+
+
+if __name__ == '__main__':
+    main()