Upload util/renderer.py with huggingface_hub

util/renderer.py

@@ -0,0 +1,49 @@
+
+import torch
+import torch.nn as nn
+import nvdiffrast.torch as dr
+from util.flexicubes_geometry import FlexiCubesGeometry
+
+class Renderer(nn.Module):
+    def __init__(self, tet_grid_size, camera_angle_num, scale, geo_type):
+        super().__init__()
+
+        self.tet_grid_size = tet_grid_size
+        self.camera_angle_num = camera_angle_num
+        self.scale = scale
+        self.geo_type = geo_type
+        # self.glctx = dr.RasterizeCudaContext()
+
+        if self.geo_type == "flex":
+            self.flexicubes = FlexiCubesGeometry(grid_res = self.tet_grid_size)   
+
+    def forward(self, data, sdf, deform, verts, tets, training=False, weight = None):
+
+        results = {}
+
+        deform = torch.tanh(deform) / self.tet_grid_size * self.scale / 0.95
+        if self.geo_type == "flex":
+            deform = deform *0.5
+
+            v_deformed = verts + deform
+
+            verts_list = []
+            faces_list = []
+            reg_list = []
+            n_shape = verts.shape[0]
+            for i in range(n_shape): 
+                verts_i, faces_i, reg_i = self.flexicubes.get_mesh(v_deformed[i], sdf[i].squeeze(dim=-1),
+                with_uv=False, indices=tets, weight_n=weight[i], is_training=training)
+
+                verts_list.append(verts_i)
+                faces_list.append(faces_i)
+                reg_list.append(reg_i)       
+            verts = verts_list
+            faces = faces_list
+
+            flexicubes_surface_reg = torch.cat(reg_list).mean()
+            flexicubes_weight_reg = (weight ** 2).mean()
+            results["flex_surf_loss"] = flexicubes_surface_reg
+            results["flex_weight_loss"] = flexicubes_weight_reg
+
+        return results, verts, faces