nerfstudio-project / nerfstudio

A collaboration friendly studio for NeRFs
https://docs.nerf.studio
Apache License 2.0
8.87k stars 1.18k forks source link

SplatFacto mask resizing bug #3199

Open ramazan793 opened 3 weeks ago

ramazan793 commented 3 weeks ago

Describe the bug I'm trying to fit a SplatFacto model on a dataset of images with masks. However, I encounter the following error:

Traceback (most recent call last):
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/bin/ns-train", line 8, in <module>
    sys.exit(entrypoint())
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/scripts/train.py", line 262, in entrypoint
    main(
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/scripts/train.py", line 247, in main
    launch(
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/scripts/train.py", line 189, in launch
    main_func(local_rank=0, world_size=world_size, config=config)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/scripts/train.py", line 100, in train_loop
    trainer.train()
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/engine/trainer.py", line 261, in train
    loss, loss_dict, metrics_dict = self.train_iteration(step)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/utils/profiler.py", line 112, in inner
    out = func(*args, **kwargs)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/engine/trainer.py", line 496, in train_iteration
    _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/utils/profiler.py", line 112, in inner
    out = func(*args, **kwargs)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/pipelines/base_pipeline.py", line 303, in get_train_loss_dict
    loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/models/splatfacto.py", line 868, in get_loss_dict
    mask = self._downscale_if_required(batch["mask"])
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/models/splatfacto.py", line 649, in _downscale_if_required
    return resize_image(image, d)
  File "/home/r.fazylov/anaconda3/envs/pytorch3d/lib/python3.10/site-packages/nerfstudio/models/splatfacto.py", line 95, in resize_image
    return tf.conv2d(image.permute(2, 0, 1)[:, None, ...], weight, stride=d).squeeze(1).permute(1, 2, 0)
RuntimeError: "slow_conv2d_cpu" not implemented for 'Bool'

It seems that the dataset converts masks to a bool tensor when loading them, while splatfacto tries to resize the mask via torch.functional.conv2d that has no implementation for boolean type.

To Reproduce Steps to reproduce the behavior: Start ns-train splatfacto on a dataset with a mask_path defined in transforms.json.

Expected behavior The training should handle masks correctly, allowing the use of masks with the objects during training.