Full model of the large reconstruction model.
| 26 | |
| 27 | |
| 28 | class InstantNeRF(nn.Module): |
| 29 | """ |
| 30 | Full model of the large reconstruction model. |
| 31 | """ |
| 32 | def __init__( |
| 33 | self, |
| 34 | encoder_freeze: bool = False, |
| 35 | encoder_model_name: str = 'facebook/dino-vitb16', |
| 36 | encoder_feat_dim: int = 768, |
| 37 | transformer_dim: int = 1024, |
| 38 | transformer_layers: int = 16, |
| 39 | transformer_heads: int = 16, |
| 40 | triplane_low_res: int = 32, |
| 41 | triplane_high_res: int = 64, |
| 42 | triplane_dim: int = 80, |
| 43 | rendering_samples_per_ray: int = 128, |
| 44 | ): |
| 45 | super().__init__() |
| 46 | |
| 47 | # modules |
| 48 | self.encoder = DinoWrapper( |
| 49 | model_name=encoder_model_name, |
| 50 | freeze=encoder_freeze, |
| 51 | ) |
| 52 | |
| 53 | self.transformer = TriplaneTransformer( |
| 54 | inner_dim=transformer_dim, |
| 55 | num_layers=transformer_layers, |
| 56 | num_heads=transformer_heads, |
| 57 | image_feat_dim=encoder_feat_dim, |
| 58 | triplane_low_res=triplane_low_res, |
| 59 | triplane_high_res=triplane_high_res, |
| 60 | triplane_dim=triplane_dim, |
| 61 | ) |
| 62 | |
| 63 | self.synthesizer = TriplaneSynthesizer( |
| 64 | triplane_dim=triplane_dim, |
| 65 | samples_per_ray=rendering_samples_per_ray, |
| 66 | ) |
| 67 | |
| 68 | def forward_planes(self, images, cameras): |
| 69 | # images: [B, V, C_img, H_img, W_img] |
| 70 | # cameras: [B, V, 16] |
| 71 | B = images.shape[0] |
| 72 | |
| 73 | # encode images |
| 74 | image_feats = self.encoder(images, cameras) |
| 75 | image_feats = rearrange(image_feats, '(b v) l d -> b (v l) d', b=B) |
| 76 | |
| 77 | # transformer generating planes |
| 78 | planes = self.transformer(image_feats) |
| 79 | |
| 80 | return planes |
| 81 | |
| 82 | def forward_synthesizer(self, planes, render_cameras, render_size: int): |
| 83 | render_results = self.synthesizer( |
| 84 | planes, |
| 85 | render_cameras, |
nothing calls this directly
no outgoing calls
no test coverage detected