Skip to content

Commit

Permalink
Add per frame intrinsic support to nerfstudio data (#1049)
Browse files Browse the repository at this point in the history
  • Loading branch information
@tancik
tancik authored Nov 29, 2022
1 parent 87c3ead commit a1f20b3
Showing 1 changed file with 77 additions and 17 deletions.
94 changes: 77 additions & 17 deletions nerfstudio/data/dataparsers/nerfstudio_dataparser.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,14 +80,64 @@ def _generate_dataparser_outputs(self, split="train"):
poses = []
num_skipped_image_filenames = 0

fx_fixed = "fl_x" in meta
fy_fixed = "fl_y" in meta
cx_fixed = "cx" in meta
cy_fixed = "cy" in meta
height_fixed = "h" in meta
width_fixed = "w" in meta
distort_fixed = False
for distort_key in ["k1", "k2", "k3", "p1", "p2"]:
if distort_key in meta:
distort_fixed = True
break
fx = []
fy = []
cx = []
cy = []
height = []
width = []
distort = []

for frame in meta["frames"]:
filepath = PurePath(frame["file_path"])
fname = self._get_fname(filepath)
if not fname.exists():
num_skipped_image_filenames += 1
else:
image_filenames.append(fname)
poses.append(np.array(frame["transform_matrix"]))
continue

if not fx_fixed:
assert "fl_x" in frame, "fx not specified in frame"
fx.append(float(frame["fl_x"]))
if not fy_fixed:
assert "fl_y" in frame, "fy not specified in frame"
fy.append(float(frame["fl_y"]))
if not cx_fixed:
assert "cx" in frame, "cx not specified in frame"
cx.append(float(frame["cx"]))
if not cy_fixed:
assert "cy" in frame, "cy not specified in frame"
cy.append(float(frame["cy"]))
if not height_fixed:
assert "h" in frame, "height not specified in frame"
height.append(int(frame["h"]))
if not width_fixed:
assert "w" in frame, "width not specified in frame"
width.append(int(frame["w"]))
if not distort_fixed:
distort.append(
camera_utils.get_distortion_params(
k1=float(meta["k1"]) if "k1" in meta else 0.0,
k2=float(meta["k2"]) if "k2" in meta else 0.0,
k3=float(meta["k3"]) if "k3" in meta else 0.0,
k4=float(meta["k4"]) if "k4" in meta else 0.0,
p1=float(meta["p1"]) if "p1" in meta else 0.0,
p2=float(meta["p2"]) if "p2" in meta else 0.0,
)
)

image_filenames.append(fname)
poses.append(np.array(frame["transform_matrix"]))
if "mask_path" in frame:
mask_filepath = PurePath(frame["mask_path"])
mask_fname = self._get_fname(mask_filepath, downsample_folder_prefix="masks_")
Expand Down Expand Up @@ -162,23 +212,33 @@ def _generate_dataparser_outputs(self, split="train"):
else:
camera_type = CameraType.PERSPECTIVE

distortion_params = camera_utils.get_distortion_params(
k1=float(meta["k1"]) if "k1" in meta else 0.0,
k2=float(meta["k2"]) if "k2" in meta else 0.0,
k3=float(meta["k3"]) if "k3" in meta else 0.0,
k4=float(meta["k4"]) if "k4" in meta else 0.0,
p1=float(meta["p1"]) if "p1" in meta else 0.0,
p2=float(meta["p2"]) if "p2" in meta else 0.0,
)
idx_tensor = torch.tensor(indices)
fx = float(meta["fl_x"]) if fx_fixed else torch.tensor(fx, dtype=torch.float32)[idx_tensor]
fy = float(meta["fl_y"]) if fy_fixed else torch.tensor(fy, dtype=torch.float32)[idx_tensor]
cx = float(meta["cx"]) if cx_fixed else torch.tensor(cx, dtype=torch.float32)[idx_tensor]
cy = float(meta["cy"]) if cy_fixed else torch.tensor(cy, dtype=torch.float32)[idx_tensor]
height = int(meta["h"]) if height_fixed else torch.tensor(height, dtype=torch.int32)[idx_tensor]
width = int(meta["w"]) if width_fixed else torch.tensor(width, dtype=torch.int32)[idx_tensor]
if distort_fixed:
distortion_params = camera_utils.get_distortion_params(
k1=float(meta["k1"]) if "k1" in meta else 0.0,
k2=float(meta["k2"]) if "k2" in meta else 0.0,
k3=float(meta["k3"]) if "k3" in meta else 0.0,
k4=float(meta["k4"]) if "k4" in meta else 0.0,
p1=float(meta["p1"]) if "p1" in meta else 0.0,
p2=float(meta["p2"]) if "p2" in meta else 0.0,
)
else:
distortion_params = torch.stack(distort, dim=0)[idx_tensor]

cameras = Cameras(
fx=float(meta["fl_x"]),
fy=float(meta["fl_y"]),
cx=float(meta["cx"]),
cy=float(meta["cy"]),
fx=fx,
fy=fy,
cx=cx,
cy=cy,
distortion_params=distortion_params,
height=int(meta["h"]),
width=int(meta["w"]),
height=height,
width=width,
camera_to_worlds=poses[:, :3, :4],
camera_type=camera_type,
)
Expand Down

2 comments on commit a1f20b3

@thekaka
Copy link

@thekaka thekaka commented on a1f20b3 Nov 30, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, @tancik , thanks for the contirbution.

I have an dataset with different intrinsic cameras, and I trained it with phototourism mode successfully.

I tried your commit with "instant-ngp" mode, and dataset is loaded correctly but this error occured afterwards.

  File "/opt/conda/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/conda/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/root/.vscode-server/extensions/ms-python.python-2022.18.2/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/__main__.py", line 39, in <module>
    cli.main()
  File "/root/.vscode-server/extensions/ms-python.python-2022.18.2/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/root/.vscode-server/extensions/ms-python.python-2022.18.2/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
    runpy.run_path(target, run_name="__main__")
  File "/root/.vscode-server/extensions/ms-python.python-2022.18.2/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 322, in run_path
    pkg_name=pkg_name, script_name=fname)
  File "/root/.vscode-server/extensions/ms-python.python-2022.18.2/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 136, in _run_module_code
    mod_name, mod_spec, pkg_name, script_name)
  File "/root/.vscode-server/extensions/ms-python.python-2022.18.2/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
    exec(code, run_globals)
  File "scripts/train.py", line 257, in <module>
    entrypoint()
  File "scripts/train.py", line 251, in entrypoint
    description=convert_markup_to_ansi(__doc__),
  File "scripts/train.py", line 240, in main
    config=config,
  File "scripts/train.py", line 173, in launch
    main_func(local_rank=0, world_size=world_size, config=config)
  File "scripts/train.py", line 87, in train_loop
    trainer.setup()
  File "/packages/nerfstudio/nerfstudio/engine/trainer.py", line 116, in setup
    device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank
  File "/packages/nerfstudio/nerfstudio/configs/base_config.py", line 66, in setup
    return self._target(self, **kwargs)
  File "/packages/nerfstudio/nerfstudio/pipelines/dynamic_batch.py", line 57, in __init__
    super().__init__(config, device, test_mode, world_size, local_rank)
  File "/packages/nerfstudio/nerfstudio/pipelines/base_pipeline.py", line 222, in __init__
    device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank
  File "/packages/nerfstudio/nerfstudio/configs/base_config.py", line 66, in setup
    return self._target(self, **kwargs)
  File "/packages/nerfstudio/nerfstudio/data/datamanagers/base_datamanager.py", line 320, in __init__
    super().__init__()
  File "/packages/nerfstudio/nerfstudio/data/datamanagers/base_datamanager.py", line 147, in __init__
    self.setup_train()
  File "/packages/nerfstudio/nerfstudio/data/datamanagers/base_datamanager.py", line 360, in setup_train
    collate_fn=self.config.collate_fn,
  File "/packages/nerfstudio/nerfstudio/data/utils/dataloaders.py", line 80, in __init__
    self.cached_collated_batch = self._get_collated_batch()
  File "/packages/nerfstudio/nerfstudio/data/utils/dataloaders.py", line 118, in _get_collated_batch
    collated_batch = self.collate_fn(batch_list)
  File "/packages/nerfstudio/nerfstudio/data/utils/nerfstudio_collate.py", line 127, in nerfstudio_collate
    {key: nerfstudio_collate([d[key] for d in batch], extra_mappings=extra_mappings) for key in elem}
  File "/packages/nerfstudio/nerfstudio/data/utils/nerfstudio_collate.py", line 127, in <dictcomp>
    {key: nerfstudio_collate([d[key] for d in batch], extra_mappings=extra_mappings) for key in elem}
  File "/packages/nerfstudio/nerfstudio/data/utils/nerfstudio_collate.py", line 107, in nerfstudio_collate
    return torch.stack(batch, 0, out=out)
RuntimeError: stack expects each tensor to be equal size, but got [1080, 1920, 3] at entry 0 and [1856, 2880, 3] at entry 1

@thekaka
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Btw, though I could train the multi-camera dataset with phototourism mode, but the evalution dosent work.

It seems that evaluation needs valdation data with the same size.

Please sign in to comment.