easyvolcap: add easier conversion from nerf dataset

scripts/nerf/nerf_to_easyvolcap.py

@@ -15,8 +15,11 @@ def main():
     parser.add_argument('--nerf_root', type=str, default='/mnt/data/home/shuaiqing/Code/MultiNBFast/cache/Hospital/4')
     parser.add_argument('--volcap_root', type=str, default='/mnt/data/home/shuaiqing/Code/MultiNBFast/cache/Hospital/4')
     parser.add_argument('--transforms_file', type=str, default='transforms_train.json')
-    parser.add_argument('--intri_file', type=str, default='intri.yml')
-    parser.add_argument('--extri_file', type=str, default='extri.yml')
+    parser.add_argument('--intri_file', type=str, default='cameras_train/00/intri.yml')
+    parser.add_argument('--extri_file', type=str, default='cameras_train/00/extri.yml')
+    parser.add_argument('--images_dir', type=str, default='images_train')
+    parser.add_argument('--no_organize_images', action='store_false', dest='organize_images')
+    parser.add_argument('--no_convert_test', action='store_false', dest='convert_test')
     args = parser.parse_args()
 
     # clean and restart
@@ -63,6 +66,12 @@ def main():
     )
     log(yellow(f'Converted cameras saved to {blue(join(args.volcap_root, f"{{{args.intri_file},{args.extri_file}}}"))}'))
 
+    if args.organize_images:
+        run(f'python scripts/nerf/organize_images.py --nerf_root {args.nerf_root} --volcap_root {args.volcap_root} --transforms_file {args.transforms_file} --images_dir {args.images_dir}')
+
+    if args.transforms_file.endswith('_train.json') and exists(join(args.nerf_root, args.transforms_file.replace("_train.json", "_test.json"))) and args.convert_test:
+        run(f'python scripts/nerf/nerf_to_easyvolcap.py --nerf_root {args.nerf_root} --volcap_root {args.volcap_root} --transforms_file {args.transforms_file.replace("_train.json", "_test.json")} --images_dir {args.images_dir.replace("_train", "_test")} --intri_file {args.intri_file.replace("_train", "_test")} --extri_file {args.extri_file.replace("_train", "_test")}')
+
 
 if __name__ == '__main__':
     main()

scripts/nerf/organize_images.py

@@ -0,0 +1,22 @@
+from easyvolcap.utils.console_utils import *
+
+@catch_throw
+def main():
+    args = dotdict()
+    args.nerf_root = '/mnt/data/home/shuaiqing/Code/MultiNBFast/cache/Hospital/4'
+    args.volcap_root = '/mnt/data/home/shuaiqing/Code/MultiNBFast/cache/Hospital/4'
+    args.images_dir = 'images_train/00'
+    args.transforms_file = 'transforms_train.json'
+    args = dotdict(vars(build_parser(args).parse_args()))
+    transforms = dotdict(json.load(open(join(args.nerf_root, args.transforms_file))))
+
+    os.makedirs(join(args.volcap_root, args.images_dir), exist_ok=True)
+    for i in tqdm(range(len(transforms.frames))):
+        src = join(args.nerf_root, transforms.frames[i].file_path)
+        tar = join(args.volcap_root, args.images_dir, f'{i:06d}.' + src.split('.')[-1])
+        if exists(tar): os.remove(tar)
+        os.symlink(relpath(src, dirname(tar)), tar)
+
+
+if __name__ == '__main__':
+    main()