facebookresearch / pytorch3d

PyTorch3D is FAIR's library of reusable components for deep learning with 3D data
https://pytorch3d.org/
Other
8.85k stars 1.32k forks source link

How can I use OpenGL-related methods in PyTorch3D to render semi-transparent textures? #1854

Open changfali opened 3 months ago

changfali commented 3 months ago

How can I use OpenGL-related methods in PyTorch3D to render semi-transparent textures? I want to achieve the effect shown in the image below, or does PyTorch3D support using the blend function like gl.glEnable(gl.GL_BLEND) in MeshRasterizerOpenGL? img_v3_02e1_f4e05203-f392-4e41-b892-bbea1021e7dg

In my scene, after rasterization, a single pixel might correspond to multiple faces. If I use MeshRasterizerOpenGL, its rasterization method is:

pix_to_face, bary_coords, zbuf = self.opengl_machinery(
    meshes_gl_ndc, projection_matrix, image_size
)

The pix_to_face result always contains only one face (with a shape of [N, H, W, 1]). However, when I use MeshRasterizer, the rasterization method is rasterize_meshes, which allows modifying the result by setting faces_per_pixel. The reason I’m not choosing MeshRasterizer is that MeshRasterizerOpenGL is faster, and I don’t care about differentiability. Is it possible to achieve my requirement using MeshRasterizerOpenGL? Or how should I modify it to get the result I want? @amyreese @JaapSuter @thatch @yurimalheiros could anyone help me with it?

changfali commented 3 months ago

I tried modifying the _set_up_gl_program_properties method in the class _OpenGLMachinery as shown below:

@staticmethod
  def _set_up_gl_program_properties(program) -> None:
      """
      Set basic OpenGL program properties: disable blending, enable depth testing,
      and disable face culling.
      """
      gl.glUseProgram(program)
      # gl.glDisable(gl.GL_BLEND)
      gl.glEnable(gl.GL_DEPTH_TEST)
      gl.glEnable(gl.GL_BLEND)     # add
      gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_CONSTANT_ALPHA)  # add
      gl.glDisable(gl.GL_CULL_FACE)
      gl.glUseProgram(0)

but I encountered the following bug:

Traceback (most recent call last):
  File "/home/lichangfa/job/Wrinkles_Synthesis/test_opengl.py", line 65, in <module>
    images = renderer(mesh_with_textures)
  File "/home/lichangfa/miniconda3/envs/PBR/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/renderer/mesh/renderer.py", line 65, in forward
    images = self.shader(fragments, meshes_world, **kwargs)
  File "/home/lichangfa/miniconda3/envs/PBR/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/renderer/mesh/shader.py", line 317, in forward
    colors = phong_shading(
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/renderer/mesh/shading.py", line 689, in phong_shading
    colors, _ = _phong_shading_with_pixels(meshes, fragments, lights, cameras, materials, texels)
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/renderer/mesh/shading.py", line 626, in _phong_shading_with_pixels
    vertex_normals = meshes.verts_normals_packed()  # (V, 3)
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/structures/meshes.py", line 727, in verts_normals_packed
    self._compute_vertex_normals()
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/structures/meshes.py", line 838, in _compute_vertex_normals
    if self.isempty():
  File "/home/lichangfa/programs/pytorch3d_aug/pytorch3d/structures/meshes.py", line 498, in isempty
    return self._N == 0 or self.valid.eq(False).all()
RuntimeError: CUDA error: an illegal memory access was encountered

here is my code:

import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import torch
import matplotlib.pyplot as plt
from pytorch3d.structures import Meshes
from pytorch3d.renderer import (
    OpenGLPerspectiveCameras,
    MeshRenderer,
    RasterizationSettings,
    PointLights,
    TexturesVertex,
    HardPhongShader,
)
from pytorch3d.renderer.opengl.rasterizer_opengl import MeshRasterizerOpenGL
from pytorch3d.utils import ico_sphere

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

mesh = ico_sphere(level=3, device=device)

verts_rgb = torch.ones_like(mesh.verts_packed())
verts_rgb[:, 0] = 1.0
verts_rgb[:, 1] = 0.0
verts_rgb[:, 2] = 0.0
textures = TexturesVertex(verts_features=verts_rgb[None])

mesh_with_textures = Meshes(verts=[mesh.verts_packed()], faces=[mesh.faces_packed()], textures=textures)

R = torch.tensor([[[1, 0, 0], [0, 1, 0], [0, 0, 1]]], device=device, dtype=torch.float32)
T = torch.tensor([[[0, 0, 3]]], device=device, dtype=torch.float32).view(1, 3)
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)

lights = PointLights(device=device, location=[[2.0, 2.0, -2.0]])

raster_settings = RasterizationSettings(
    image_size=2048,
    blur_radius=0.0,
    faces_per_pixel=1,
)

renderer = MeshRenderer(
    rasterizer=MeshRasterizerOpenGL(cameras=cameras, raster_settings=raster_settings),
    shader=HardPhongShader(
        device=device, cameras=cameras, lights=PointLights(device=device, location=cameras.get_camera_center())
    ).to(device),
)

images = renderer(mesh_with_textures)

image = images[0, ..., :3].cpu().numpy()
plt.imshow(image)
plt.grid(False)
plt.axis("off")
plt.show()
bottler commented 3 months ago

I'm afraid this is beyond what we can help with at the moment. Sounds like an interesting and useful idea.

cajoek commented 1 month ago

@changfali sorry to hijack this issue but could you please tell me how you installed pytorch3d to get MeshRasterizerOpenGL working? Thanks :pray: