update

refine/func.py

@@ -14,6 +14,37 @@ import pymeshlab as ml
 from pymeshlab import Percentage
 import nvdiffrast.torch as dr
 import numpy as np
+from pytorch3d.renderer import MeshRasterizer
+from pytorch3d.renderer.mesh.rasterizer import Fragments
+
+def render_pix2faces_py3d(meshes, cameras, H=512, W=512, blur_radius=0.0, faces_per_pixel=1):
+    """
+    Renders pix2face of visible faces.
+    :param mesh: Pytorch3d.structures.Meshes
+    :param cameras: pytorch3d.renderer.Cameras
+    :param H: target image height
+    :param W: target image width
+    :param blur_radius: Float distance in the range [0, 2] used to expand the face
+            bounding boxes for rasterization. Setting blur radius
+            results in blurred edges around the shape instead of a
+            hard boundary. Set to 0 for no blur.
+    :param faces_per_pixel: (int) Number of faces to keep track of per pixel.
+            We return the nearest faces_per_pixel faces along the z-axis.
+    """
+    # Define the settings for rasterization and shading
+    raster_settings = RasterizationSettings(
+        image_size=(H, W),
+        blur_radius=blur_radius,
+        faces_per_pixel=faces_per_pixel
+    )
+    rasterizer=MeshRasterizer(
+        cameras=cameras, 
+        raster_settings=raster_settings
+    )
+    fragments: Fragments = rasterizer(meshes, cameras=cameras)
+    return {
+        "pix_to_face": fragments.pix_to_face[..., 0],
+    }
 
 
 def _translation(x, y, z, device):
@@ -162,11 +193,12 @@ class Pix2FacesRenderer:
         pix_to_face = rast_out[..., -1].to(torch.int32) - 1
         return pix_to_face
 
-pix2faces_renderer = Pix2FacesRenderer()
+# pix2faces_renderer = Pix2FacesRenderer()
+pix2faces_renderer = None
 
 def get_visible_faces(meshes: Meshes, cameras: CamerasBase, resolution=1024):
-    # pix_to_face = render_pix2faces_py3d(meshes, cameras, H=resolution, W=resolution)['pix_to_face']
-    pix_to_face = pix2faces_renderer.render_pix2faces_nvdiff(meshes, cameras, H=resolution, W=resolution)
+    pix_to_face = render_pix2faces_py3d(meshes, cameras, H=resolution, W=resolution)['pix_to_face']
+    # pix_to_face = pix2faces_renderer.render_pix2faces_nvdiff(meshes, cameras, H=resolution, W=resolution)
 
     unique_faces = torch.unique(pix_to_face.flatten())
     unique_faces = unique_faces[unique_faces != -1]

refine/render.py

@@ -15,7 +15,7 @@ def _warmup(glctx, device=None):
     tri = tensor([[0, 1, 2]], dtype=torch.int32)
     dr.rasterize(glctx, pos, tri, resolution=[256, 256])
 
-glctx = dr.RasterizeCudaContext(device="cuda")
+glctx = dr.RasterizeCudaContext()
 
 class NormalsRenderer: