125 Commits

Author SHA1 Message Date
Jeremy Reizenstein
83bacda8fb lint
Summary: Fix recent flake complaints

Reviewed By: MichaelRamamonjisoa

Differential Revision: D51811912

fbshipit-source-id: 65183f5bc7058da910e4d5a63b2250ce8637f1cc
2023-12-04 13:43:34 -08:00
Emilien Garreau
9446d91fae Avoid to keep in memory lengths and bins for ImplicitronRayBundle
Summary:
Convert ImplicitronRayBundle to a "classic" class instead of a dataclass. This change is introduced as a way to preserve the ImplicitronRayBundle interface while allowing two outcomes:
- init lengths arguments is now a Optional[torch.Tensor] instead of torch.Tensor
- lengths is now a property which returns a `torch.Tensor`. The lengths property will either recompute lengths from bins or return the stored _lengths. `_lenghts` is None if bins is set. It saves us a bit of memory.

Reviewed By: shapovalov

Differential Revision: D46686094

fbshipit-source-id: 3c75c0947216476ebff542b6f552d311024a679b
2023-07-06 02:41:15 -07:00
Emilien Garreau
3d011a9198 Adapt RayPointRefiner and RayMarcher to support bins.
Summary:
## Context

Bins are used in mipnerf to allow to manipulate easily intervals. For example, by doing the following, `bins[..., :-1]` you will obtain all the left coordinates of your intervals, while doing `bins[..., 1:]` is equals to the right coordinates of your intervals.

We introduce here the support of bins like in MipNerf implementation.

## RayPointRefiner

Small changes have been made to modify RayPointRefiner.
- If bins is None

```
mids = torch.lerp(ray_bundle.lengths[..., 1:], ray_bundle.lengths[…, :-1], 0.5)
z_samples = sample_pdf(
		mids, # [..., npt]
		weights[..., 1:-1], # [..., npt - 1]
               ….
            )
```

- If bins is not None
In the MipNerf implementation the sampling is done on all the bins. It allows us to use the full weights tensor without slashing it.

```
z_samples = sample_pdf(
		ray_bundle.bins, # [..., npt + 1]
		weights, # [..., npt]
               ...
            )
```

## RayMarcher

Add a ray_deltas optional argument. If None, keep the same deltas computation from ray_lengths.

Reviewed By: shapovalov

Differential Revision: D46389092

fbshipit-source-id: d4f1963310065bd31c1c7fac1adfe11cbeaba606
2023-07-06 02:41:15 -07:00
Emilien Garreau
5910d81b7b Add blurpool following MIPNerf paper.
Summary:
Add blurpool has defined in [MIP-NeRF](https://arxiv.org/abs/2103.13415).
It has been added has an option for RayPointRefiner.

Reviewed By: shapovalov

Differential Revision: D46356189

fbshipit-source-id: ad841bad86d2b591a68e1cb885d4f781cf26c111
2023-07-06 02:20:53 -07:00
Emilien Garreau
ccf860f1db Add integrated position encoding based on MIPNerf implementation.
Summary: Add a new implicit module Integral Position Encoding based on [MIP-NeRF](https://arxiv.org/abs/2103.13415).

Reviewed By: shapovalov

Differential Revision: D46352730

fbshipit-source-id: c6a56134c975d80052b3a11f5e92fd7d95cbff1e
2023-07-06 02:20:53 -07:00
Emilien Garreau
29b8ebd802 Add utils to approximate the conical frustums as multivariate gaussians.
Summary:
Introduce methods to approximate the radii of conical frustums along rays as described in [MipNerf](https://arxiv.org/abs/2103.13415):

- Two new attributes are added to ImplicitronRayBundle: bins and radii. Bins is of size n_pts_per_ray + 1. It allows us to manipulate easily and n_pts_per_ray intervals. For example we need the intervals coordinates in the radii computation for \(t_{\mu}, t_{\delta}\). Radii are used to store the radii of the conical frustums.

- Add 3 new methods to compute the radii:
   - approximate_conical_frustum_as_gaussians: It computes the mean along the ray direction, the variance of the
      conical frustum  with respect to t and variance of the conical frustum with respect to its radius. This
      implementation follows the stable computation defined in the paper.
   - compute_3d_diagonal_covariance_gaussian: Will leverage the two previously computed variances to find the
     diagonal covariance of the Gaussian.
   - conical_frustum_to_gaussian: Mix everything together to compute the means and the diagonal covariances along
     the ray of the Gaussians.

- In AbstractMaskRaySampler, introduces the attribute `cast_ray_bundle_as_cone`. If False it won't change the previous behaviour of the RaySampler. However if True, the samplers will sample `n_pts_per_ray +1` instead of `n_pts_per_ray`. This points are then used to set the bins attribute of ImplicitronRayBundle. The support of HeterogeneousRayBundle has not been added since the current code does not allow it. A safeguard has been added to avoid a silent bug in the future.

Reviewed By: shapovalov

Differential Revision: D45269190

fbshipit-source-id: bf22fad12d71d55392f054e3f680013aa0d59b78
2023-07-06 01:55:41 -07:00
Jeremy Reizenstein
42e7de418c fix test_build internal
Summary: Make test work in isolation, and when run internally make it not try the sqlalchemy files.

Reviewed By: shapovalov

Differential Revision: D46352513

fbshipit-source-id: 7417a25d7a5347d937631c9f56ae4e3242dd622e
2023-06-16 04:49:02 -07:00
Roman Shapovalov
cd5db076d5 Adding SQL dataset classes to ImplicitronDataSource imports
Summary: Making it easier for the clients to use these datasets.

Reviewed By: bottler

Differential Revision: D46727179

fbshipit-source-id: cf619aee4c4c0222a74b30ea590cf37f08f014cc
2023-06-14 10:51:47 -07:00
Roman Shapovalov
09a99f2e6d Support limiting num sequences per category.
Summary:
Adds stratified sampling of sequences within categories applied after category / sequence filters but before the num sequence limit.
It respects the insertion order into the sequence_annots table, i.e. takes top N sequences within each category.

Reviewed By: bottler

Differential Revision: D46724002

fbshipit-source-id: 597cb2a795c3f3bc07f838fc51b4e95a4f981ad3
2023-06-14 07:12:02 -07:00
Emilien Garreau
35badc0892 Fix inversion between fine and coarse implicit_functions
Summary: Fine implicit function was called before the coarse implicit function.

Reviewed By: shapovalov

Differential Revision: D46224224

fbshipit-source-id: 6b1cc00cc823d3ea7a5b42774c9ec3b73a69edb5
2023-05-26 08:33:54 -07:00
Roman Shapovalov
d2119c285f Serialising dynamic arrays in SQL; read-only SQLite connection in SQL Dataset
Summary:
1. We may need to store arrays of unknown shape in the database. It implements and tests serialisation.

2. Previously, when an inexisting metadata file was passed to SqlIndexDataset, it would try to open it and create an empty file, then crash. We now open the file in a read-only mode, so the error message is more intuitive. Note that the implementation is SQLite specific.

Reviewed By: bottler

Differential Revision: D46047857

fbshipit-source-id: 3064ae4f8122b4fc24ad3d6ab696572ebe8d0c26
2023-05-22 02:24:49 -07:00
Jeremy Reizenstein
ff80183fdb resources fix
Summary: I don't know why RE tests sometimes fail here, but maybe it's a race condition. If that's right, this should fix it.

Reviewed By: shapovalov

Differential Revision: D46020054

fbshipit-source-id: 20b746b09ad9bd77c2601ac681047ccc6cc27ed9
2023-05-19 06:41:33 -07:00
Roman Shapovalov
b0462598ac Refactor: FrameDataBuilder is more extensible.
Summary:
This is mostly a refactoring diff to reduce friction in extending the frame data.

Slight functional changes: dataset getitem now accepts (seq_name, frame_number_as_singleton_tensor) as a non-advertised feature. Otherwise this code crashes:
```
item = dataset[0]
dataset[item.sequence_name, item.frame_number]
```

Reviewed By: bottler

Differential Revision: D45780175

fbshipit-source-id: 75b8e8d3dabed954a804310abdbd8ab44a8dea29
2023-05-17 10:38:34 -07:00
Roman Shapovalov
c8d6cd427e Fix test_data_source in OSS.
Summary: Import generic path; avoiding incorrect path patching.

Reviewed By: bottler

Differential Revision: D45573976

fbshipit-source-id: e6ff4d759deb936e3b636defa1e0851fb0127b46
2023-05-05 02:05:50 -07:00
Roman Shapovalov
3e3644e534 More tests for SQL Dataset
Summary:
I forgot to include these tests to D45086611 when transferring code from pixar_replay repo.

They test the new ORM types used in SQL dataset and are SQL Alchemy 2.0 specific.

An important test for extending types is a proof of concept for generality of SQL Dataset. The idea is to extend FrameAnnotation and FrameData in parallel.

Reviewed By: bottler

Differential Revision: D45529284

fbshipit-source-id: 2a634e518f580c312602107c85fc320db43abcf5
2023-05-04 03:32:27 -07:00
Roman Shapovalov
32e1992924 SQL Index Dataset
Summary:
Moving SQL dataset to PyTorch3D. It has been extensively tested in pixar_replay.

It requires SQLAlchemy 2.0, which is not supported in fbcode. So I exclude the sources and tests that depend on it from buck TARGETS.

Reviewed By: bottler

Differential Revision: D45086611

fbshipit-source-id: 0285f03e5824c0478c70ad13731525bb5ec7deef
2023-04-25 09:56:15 -07:00
Roman Shapovalov
7aeedd17a4 When bounding boxes are cached in metadata, don’t crash on load_masks=False
Summary:
We currently support caching bounding boxes in MaskAnnotation. If present, they are not re-computed from the mask. However, the masks need to be loaded for the bbox to be set.

This diff fixes that. Even if load_masks / load_blobs are unset, the bounding box can be picked up from the metadata.

Reviewed By: bottler

Differential Revision: D45144918

fbshipit-source-id: 8a2e2c115e96070b6fcdc29cbe57e1cee606ddcd
2023-04-20 07:28:45 -07:00
Roman Shapovalov
0e3138eca8 Optional ground-truth depth maps in visualiser
Summary: The code does not crash if depth map/mask are not given.

Reviewed By: bottler

Differential Revision: D45082985

fbshipit-source-id: 3610d8beb4ac897fbbe52f56a6dd012a6365b89b
2023-04-18 07:00:17 -07:00
Ildar Salakhiev
ebdbfde0ce Extract BlobLoader class from JsonIndexDataset and moving crop_by_bbox to FrameData
Summary:
extracted blob loader
added documentation for blob_loader
did some refactoring on fields
for detailed steps and discussions see:
https://github.com/facebookresearch/pytorch3d/pull/1463
https://github.com/fairinternal/pixar_replay/pull/160

Reviewed By: bottler

Differential Revision: D44061728

fbshipit-source-id: eefb21e9679003045d73729f96e6a93a1d4d2d51
2023-04-04 07:17:43 -07:00
Dejan Kovachev
c759fc560f Hard population of registry system with pre_expand
Summary: Provide an extension point pre_expand to let a configurable class A make sure another class B is registered before A is expanded. This reduces top level imports.

Reviewed By: bottler

Differential Revision: D44504122

fbshipit-source-id: c418bebbe6d33862d239be592d9751378eee3a62
2023-03-31 07:44:38 -07:00
Emilien Garreau
813e941de5 Add the OverfitModel
Summary:
Introduces the OverfitModel for NeRF-style training with overfitting to one scene.
It is a specific case of GenericModel. It has been disentangle to ease usage.

## General modification

1. Modularize a minimum GenericModel to introduce OverfitModel
2. Introduce OverfitModel and ensure through unit testing that it behaves like GenericModel.

## Modularization

The following methods have been extracted from GenericModel to allow modularity with ManyViewModel:
- get_objective is now a call to weighted_sum_losses
- log_loss_weights
- prepare_inputs

The generic methods have been moved to an utils.py file.

Simplify the code to introduce OverfitModel.

Private methods like chunk_generator are now public and can now be used by ManyViewModel.

Reviewed By: shapovalov

Differential Revision: D43771992

fbshipit-source-id: 6102aeb21c7fdd56aa2ff9cd1dd23fd9fbf26315
2023-03-24 07:27:39 -07:00
Jeremy Reizenstein
97f8f9bf47 test fix
Reviewed By: shapovalov

Differential Revision: D42780711

fbshipit-source-id: 075fcae5097147b782f7ffc935f5430b824f58fd
2023-01-27 03:56:53 -08:00
Roman Shapovalov
11959e0b24 Subsets in dataset iterators
Summary: For the new API, filtering iterators over sequences by subsets is quite helpful. The change is backwards compatible.

Reviewed By: bottler

Differential Revision: D42739669

fbshipit-source-id: d150a404aeaf42fd04a81304c63a4cba203f897d
2023-01-26 03:00:46 -08:00
Jeremy Reizenstein
a12612a48f doc rgbd point cloud
Summary: docstring and shape fix

Reviewed By: shapovalov

Differential Revision: D42609661

fbshipit-source-id: fd50234872ad61b5452821eeb89d51344f70c957
2023-01-24 15:26:52 -08:00
Jeremy Reizenstein
dcced4fa29 test fixes
Reviewed By: shapovalov

Differential Revision: D42545069

fbshipit-source-id: e25fb4049dcebd715df43bab3ce813ecb5f85abe
2023-01-17 06:11:56 -08:00
Jeremy Reizenstein
60ab1cdb72 make x_enabled compulsory
Summary: Optional[some_configurable] won't autogenerate the enabled flag

Reviewed By: shapovalov

Differential Revision: D41522104

fbshipit-source-id: 555ff6b343faf6f18aad2f92fbb7c341f5e991c6
2022-11-24 09:38:02 -08:00
Jeremy Reizenstein
a1f2ded58a voxel_grid_implicit_function scaffold fixes
Summary: Fix indexing of directions after filtering of points by scaffold.

Reviewed By: shapovalov

Differential Revision: D40853482

fbshipit-source-id: 9cfdb981e97cb82edcd27632c5848537ed2c6837
2022-11-03 05:46:31 -07:00
David Novotny
e4a3298149 CO3Dv2 multi-category extension
Summary:
Allows loading of multiple categories.
Multiple categories are provided in a comma-separated list of category names.

Reviewed By: bottler, shapovalov

Differential Revision: D40803297

fbshipit-source-id: 863938be3aa6ffefe9e563aede4a2e9e66aeeaa8
2022-11-02 13:55:25 -07:00
Jeremy Reizenstein
74754bbf17 voxel_grid_implicit_function
Reviewed By: shapovalov

Differential Revision: D40622304

fbshipit-source-id: 277515a55c46d9b8300058b439526539a7fe00a0
2022-10-23 05:36:34 -07:00
Jeremy Reizenstein
fe5bdb2fb5 different learning rate for different parts
Summary:
Adds the ability to have different learning rates for different parts of the model. The trainable parts of the implicitron have a new member

       param_groups: dictionary where keys are names of individual parameters,
            or module’s members and values are the parameter group where the
            parameter/member will be sorted to. "self" key is used to denote the
            parameter group at the module level. Possible keys, including the "self" key
            do not have to be defined. By default all parameters are put into "default"
            parameter group and have the learning rate defined in the optimizer,
            it can be overriden at the:
                - module level with “self” key, all the parameters and child
                    module s parameters will be put to that parameter group
                - member level, which is the same as if the `param_groups` in that
                    member has key=“self” and value equal to that parameter group.
                    This is useful if members do not have `param_groups`, for
                    example torch.nn.Linear.
                - parameter level, parameter with the same name as the key
                    will be put to that parameter group.

And in the optimizer factory, parameters and their learning rates are recursively gathered.

Reviewed By: shapovalov

Differential Revision: D40145802

fbshipit-source-id: 631c02b8d79ee1c0eb4c31e6e42dbd3d2882078a
2022-10-18 15:58:18 -07:00
Jeremy Reizenstein
4d9215b3b4 fix to get_default_args(instance)
Summary:
Small config system fix. Allows get_default_args to work on an instance which has been created with a dict (instead of a DictConfig) as an args field. E.g.

```
gm = GenericModel(
        raysampler_AdaptiveRaySampler_args={"scene_extent": 4.0}
    )
    OmegaConf.structured(gm1)
```

Reviewed By: shapovalov

Differential Revision: D40341047

fbshipit-source-id: 587d0e8262e271df442a80858949a48e5d6db3df
2022-10-13 06:05:07 -07:00
Darijan Gudelj
76cddd90be Elementwise decoder
Summary: Tensorf does relu or softmax after the density grid. This diff adds the ability to replicate that.

Reviewed By: bottler

Differential Revision: D40023228

fbshipit-source-id: 9f19868cd68460af98ab6e61c7f708158c26dc08
2022-10-13 05:59:22 -07:00
Darijan Gudelj
f55d37f07d volume cropping
Summary:
TensoRF at step 2000 does volume croping and resizing.
At those steps it calculates part of the voxel grid which has density big enough to have objects and resizes the grid to fit that object.
Change is done on 3 levels:
- implicit function subscribes to epochs and at specific epochs finds the bounding box of the object and calls resizing of the color and density voxel grids to fit it
- VoxelGrid module calls cropping of the underlaying voxel grid and resizing to fit previous size it also adjusts its extends and translation to match wanted size
- Each voxel grid has its own way of cropping the underlaying data

Reviewed By: kjchalup

Differential Revision: D39854548

fbshipit-source-id: 5435b6e599aef1eaab980f5421d3369ee4829c50
2022-10-12 08:31:51 -07:00
Darijan Gudelj
56d3465b09 scaffold
Summary: Forward method is sped up using the scaffold, a low resolution voxel grid which is used to filter out the points in empty space. These points will be predicted as having 0 density and (0, 0, 0) color. The points which were not evaluated as empty space will be passed through the steps outlined above.

Reviewed By: kjchalup

Differential Revision: D39579671

fbshipit-source-id: 8eab8bb43ef77c2a73557efdb725e99a6c60d415
2022-10-10 11:01:00 -07:00
Jeremy Reizenstein
95a2acf763 Co3Dv2 point cloud fix
Summary: Avoid certain hardcoded paths in co3dv2 data

Reviewed By: davnov134

Differential Revision: D40209309

fbshipit-source-id: 0e83a15baa47d5bd07d2d23c6048cb4522c1ccba
2022-10-09 05:06:49 -07:00
Darijan Gudelj
c311a4cbb9 Enable mixed frame raysampling
Summary:
Changed ray_sampler and metrics to be able to use mixed frame raysampling.

Ray_sampler now has a new member which it passes to the pytorch3d raysampler.
If the raybundle is heterogeneous metrics now samples images by padding xys first. This reduces memory consumption.

Reviewed By: bottler, kjchalup

Differential Revision: D39542221

fbshipit-source-id: a6fec23838d3049ae5c2fd2e1f641c46c7c927e3
2022-10-03 08:36:47 -07:00
Darijan Gudelj
ad8907d373 ImplicitronRayBundle
Summary: new implicitronRayBundle with added cameraIDs and camera counts. Added to enable a single raybundle inside Implicitron and easier extension in the future. Since RayBundle is named tuple and RayBundleHeterogeneous is dataclass and RayBundleHeterogeneous cannot inherit RayBundle. So if there was no ImplicitronRayBundle every function that uses RayBundle now would have to use Union[RayBundle, RaybundleHeterogeneous] which is confusing and unecessary complicated.

Reviewed By: bottler, kjchalup

Differential Revision: D39262999

fbshipit-source-id: ece160e32f6c88c3977e408e966789bf8307af59
2022-10-03 08:36:47 -07:00
Roman Shapovalov
74bbd6fd76 Fix returning a proper rotation in levelling; supporting batches and default centroid
Summary:
`get_rotation_to_best_fit_xy` is useful to expose externally, however there was a bug (which we probably did not care about for our use case): it could return a rotation matrix with det(R) == −1.
The diff fixes that, and also makes centroid optional (it can be computed from points).

Reviewed By: bottler

Differential Revision: D39926791

fbshipit-source-id: 5120c7892815b829f3ddcc23e93d4a5ec0ca0013
2022-09-29 11:56:14 -07:00
Darijan Gudelj
5005f09118 Add rescaling to voxel grids
Summary: Any module can be subscribed to step updates from the training loop. Once the training loop publishes a step the voxel grid changes its dimensions. During the construction of VoxelGridModule and its parameters it does not know which is the resolution that will be loaded from checkpoint, so before the checkpoint loading a hook runs which changes the VoxelGridModule's parameters to match shapes of the loaded checkpoint.

Reviewed By: bottler

Differential Revision: D39026775

fbshipit-source-id: 0d359ea5c8d2eda11d773d79c7513c83585d5f17
2022-09-28 05:23:22 -07:00
Jeremy Reizenstein
d6a197be36 make expand_args_fields optional
Summary: Call expand_args_field when instantiating an object.

Reviewed By: shapovalov

Differential Revision: D39541931

fbshipit-source-id: de8e1038927ff0112463394412d5d8c26c4a1e17
2022-09-22 08:36:09 -07:00
Darijan Gudelj
db3c12abfb arbitrary shape input to voxel_grids
Summary: Add the ability to process arbitrary point shapes `[n_grids, ..., 3]` instead of  only `[n_grids, n_points, 3]`.

Reviewed By: bottler

Differential Revision: D39574373

fbshipit-source-id: 0a9ecafe9ea58cd8f909644de43a1185ecf934f4
2022-09-22 03:35:11 -07:00
Jeremy Reizenstein
da7fe2854e small fixes to config
Summary:
- indicate location of OmegaConf.structured failures
- split the data gathering from enable_get_default_args to ease experimenting with it.
- comment fixes.
- nicer error when a_class_type has weird type.

Reviewed By: kjchalup

Differential Revision: D39434447

fbshipit-source-id: b80c7941547ca450e848038ef5be95b7ebbe8f3e
2022-09-15 03:03:52 -07:00
David Novotny
f6d43eaa62 Bugfixes in render_flyaround
Summary: Fixes bugs in render_flyaround

Reviewed By: bottler

Differential Revision: D39271932

fbshipit-source-id: 07e6c9ee07ba91feb437b725af0a8942fd98db0b
2022-09-06 05:45:31 -07:00
David Novotny
c79c954dea Rename and move render_flyaround into core implicitron
Summary:
Move the flyaround rendering function into core implicitron.
The unblocks an example in the facebookresearch/co3d repo.

Reviewed By: bottler

Differential Revision: D39257801

fbshipit-source-id: 6841a88a43d4aa364dd86ba83ca2d4c3cf0435a4
2022-09-06 03:03:08 -07:00
Darijan Gudelj
72c3a0ebe5 raybundle input to ImplicitFunctions -> api unification
Summary: Currently some implicit functions in implicitron take a raybundle, others take ray_points_world. raybundle is what they really need. However, the raybundle is going to become a bit more flexible later, as it will contain different numbers of rays for each camera.

Reviewed By: bottler

Differential Revision: D39173751

fbshipit-source-id: ebc038e426d22e831e67a18ba64655d8a61e1eb9
2022-09-05 06:26:06 -07:00
Darijan Gudelj
e7c609f198 Decoding functions
Summary: Added replacable decoding functions which will be applied after the voxel grid to get color and density

Reviewed By: bottler

Differential Revision: D38829763

fbshipit-source-id: f21ce206c1c19548206ea2ce97d7ebea3de30a23
2022-08-26 08:47:30 -07:00
Darijan Gudelj
24f5f4a3e7 VoxelGridModule
Summary: Simple wrapper around voxel grids to make them a module

Reviewed By: bottler

Differential Revision: D38829762

fbshipit-source-id: dfee85088fa3c65e396cc7d3bf7ebaaffaadb646
2022-08-25 09:40:43 -07:00
Darijan Gudelj
edee25a1e5 voxel grids with interpolation
Summary: Added voxel grid classes from TensoRF, both in their factorized (CP and VM) and full form.

Reviewed By: bottler

Differential Revision: D38465419

fbshipit-source-id: 8b306338af58dc50ef47a682616022a0512c0047
2022-08-23 07:22:41 -07:00
Jeremy Reizenstein
fdaaa299a7 remove stray "generic_model_args" references
Summary:
generic_model_args no longer exists. Update some references to it, mostly in doc.

This fixes the testing of all the yaml files in test_forward pass.

Reviewed By: shapovalov

Differential Revision: D38789202

fbshipit-source-id: f11417efe772d7f86368b3598aa66c52b1309dbf
2022-08-18 07:18:55 -07:00
Jeremy Reizenstein
d42e0d3d86 1GPU for implicitron tests
Reviewed By: shapovalov

Differential Revision: D38794764

fbshipit-source-id: 140c8a935d760bab8569d903cc52ac3dd73cd553
2022-08-18 03:23:21 -07:00