chores: rebase commits

threestudio/models/renderers/patch_renderer.py

@@ -0,0 +1,106 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.renderers.base import VolumeRenderer
+from threestudio.utils.typing import *
+
+
+@threestudio.register("patch-renderer")
+class PatchRenderer(VolumeRenderer):
+    @dataclass
+    class Config(VolumeRenderer.Config):
+        patch_size: int = 128
+        base_renderer_type: str = ""
+        base_renderer: Optional[VolumeRenderer.Config] = None
+        global_detach: bool = False
+        global_downsample: int = 4
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        self.base_renderer = threestudio.find(self.cfg.base_renderer_type)(
+            self.cfg.base_renderer,
+            geometry=geometry,
+            material=material,
+            background=background,
+        )
+
+    def forward(
+        self,
+        rays_o: Float[Tensor, "B H W 3"],
+        rays_d: Float[Tensor, "B H W 3"],
+        light_positions: Float[Tensor, "B 3"],
+        bg_color: Optional[Tensor] = None,
+        **kwargs
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        B, H, W, _ = rays_o.shape
+
+        if self.base_renderer.training:
+            downsample = self.cfg.global_downsample
+            global_rays_o = torch.nn.functional.interpolate(
+                rays_o.permute(0, 3, 1, 2),
+                (H // downsample, W // downsample),
+                mode="bilinear",
+            ).permute(0, 2, 3, 1)
+            global_rays_d = torch.nn.functional.interpolate(
+                rays_d.permute(0, 3, 1, 2),
+                (H // downsample, W // downsample),
+                mode="bilinear",
+            ).permute(0, 2, 3, 1)
+            out_global = self.base_renderer(
+                global_rays_o, global_rays_d, light_positions, bg_color, **kwargs
+            )
+
+            PS = self.cfg.patch_size
+            patch_x = torch.randint(0, W - PS, (1,)).item()
+            patch_y = torch.randint(0, H - PS, (1,)).item()
+            patch_rays_o = rays_o[:, patch_y : patch_y + PS, patch_x : patch_x + PS]
+            patch_rays_d = rays_d[:, patch_y : patch_y + PS, patch_x : patch_x + PS]
+            out = self.base_renderer(
+                patch_rays_o, patch_rays_d, light_positions, bg_color, **kwargs
+            )
+
+            valid_patch_key = []
+            for key in out:
+                if torch.is_tensor(out[key]):
+                    if len(out[key].shape) == len(out["comp_rgb"].shape):
+                        if out[key][..., 0].shape == out["comp_rgb"][..., 0].shape:
+                            valid_patch_key.append(key)
+            for key in valid_patch_key:
+                out_global[key] = F.interpolate(
+                    out_global[key].permute(0, 3, 1, 2), (H, W), mode="bilinear"
+                ).permute(0, 2, 3, 1)
+                if self.cfg.global_detach:
+                    out_global[key] = out_global[key].detach()
+                out_global[key][
+                    :, patch_y : patch_y + PS, patch_x : patch_x + PS
+                ] = out[key]
+            out = out_global
+        else:
+            out = self.base_renderer(
+                rays_o, rays_d, light_positions, bg_color, **kwargs
+            )
+
+        return out
+
+    def update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ) -> None:
+        self.base_renderer.update_step(epoch, global_step, on_load_weights)
+
+    def train(self, mode=True):
+        return self.base_renderer.train(mode)
+
+    def eval(self):
+        return self.base_renderer.eval()