new

shader.py

@@ -0,0 +1,76 @@
+import torch
+from pytorch3d.renderer.mesh.shader import ShaderBase
+from pytorch3d.renderer import (
+    SoftPhongShader,
+)
+
+class MultiOutputShader(ShaderBase):
+    def __init__(self, device, cameras, lights, materials, ccm_scale=1.0, choices=None):
+        super().__init__()
+        self.device = device
+        self.cameras = cameras
+        self.lights = lights
+        self.materials = materials
+        self.ccm_scale = ccm_scale
+
+        if choices is None:
+            self.choices = ["rgb", "mask", "depth", "normal", "albedo", "ccm"]
+        else:
+            self.choices = choices
+
+        self.phong_shader = SoftPhongShader(
+            device=self.device,
+            cameras=self.cameras,
+            lights=self.lights,
+            materials=self.materials
+        )
+
+    def forward(self, fragments, meshes, **kwargs):
+        batch_size, H, W, _ = fragments.zbuf.shape
+        output = {}
+
+        if "rgb" in self.choices:
+            rgb_images = self.phong_shader(fragments, meshes, **kwargs)
+            rgb = rgb_images[..., :3]
+            output["rgb"] = rgb
+        
+        if "mask" in self.choices:
+            alpha = rgb_images[..., 3:4]
+            mask = (alpha > 0).float()
+            output["mask"] = mask
+        
+        if "albedo" in self.choices:
+            albedo = meshes.sample_textures(fragments)
+            output["albedo"] = albedo[..., 0, :]
+        
+        if "depth" in self.choices:
+            depth = fragments.zbuf
+            output["depth"] = depth
+
+        if "normal" in self.choices:
+            pix_to_face = fragments.pix_to_face[..., 0]
+            bary_coords = fragments.bary_coords[..., 0, :]
+            valid_mask = pix_to_face >= 0
+            face_indices = pix_to_face[valid_mask]
+            faces_packed = meshes.faces_packed()
+            normals_packed = meshes.verts_normals_packed()
+            face_vertex_normals = normals_packed[faces_packed[face_indices]] 
+            bary = bary_coords.view(-1, 3)[valid_mask.view(-1)]
+            interpolated_normals = (
+                bary[..., 0:1] * face_vertex_normals[:, 0, :] +
+                bary[..., 1:2] * face_vertex_normals[:, 1, :] +
+                bary[..., 2:3] * face_vertex_normals[:, 2, :]
+            )
+            interpolated_normals = interpolated_normals / interpolated_normals.norm(dim=-1, keepdim=True)
+            normal = torch.zeros(batch_size, H, W, 3, device=self.device)
+            normal[valid_mask] = interpolated_normals
+            output["normal"] = normal
+
+        if "ccm" in self.choices:
+            face_vertices = meshes.verts_packed()[meshes.faces_packed()]
+            faces_at_pixels = face_vertices[fragments.pix_to_face]
+            ccm = torch.sum(fragments.bary_coords.unsqueeze(-1) * faces_at_pixels, dim=-2)
+            ccm = (ccm[..., 0, :] * self.ccm_scale + 1) / 2
+            output["ccm"] = ccm
+
+        return output

video_render.py

@@ -1,28 +1,28 @@
-import os
-import math
-import numpy as np
+import pytorch3d
+import torch
 import imageio
+import numpy as np
+import os
+from pytorch3d.io import load_objs_as_meshes
+from pytorch3d.renderer import (
+    AmbientLights,
+    PerspectiveCameras, 
+    RasterizationSettings, 
+    look_at_view_transform,
+    TexturesVertex,
+    MeshRenderer, 
+    Materials,
+    MeshRasterizer, 
+    SoftPhongShader, 
+    PointLights
+)
 import trimesh
-import pyrender
 from tqdm import tqdm
-# os.environ["CUDA_VISIBLE_DEVICES"] = "7"
-os.environ['PYOPENGL_PLATFORM'] = 'egl'  # 设置渲染环境为 EGL（无头模式）
-
-def render_video_from_obj(input_obj_path, output_video_path, fps=15, frame_count=60, resolution=(512, 512)):
-    """
-    Render a rotating 3D model (OBJ file) to a video with RGB and normal map side-by-side.
-    
-    Args:
-        input_obj_path (str): Path to the input OBJ file.
-        output_video_path (str): Path to save the output video.
-        fps (int): Frames per second for the video.
-        frame_count (int): Number of frames in the video.
-        resolution (tuple): Resolution of the rendered video (width, height).
-        
-    Returns:
-        str: Path to the output video.
-    """
-    # 检查输入文件是否存在
+from pytorch3d.transforms import RotateAxisAngle
+
+from shader import MultiOutputShader
+
+def render_video_from_obj(obj_path, output_video_path, num_frames=60, image_size=512, fps=30, device="cuda"):
     if not os.path.exists(input_obj_path):
         raise FileNotFoundError(f"Input OBJ file not found: {input_obj_path}")
 
@@ -39,85 +39,98 @@ def render_video_from_obj(input_obj_path, output_video_path, fps=15, frame_count
     if not hasattr(mesh_data, 'vertex_normals') or mesh_data.vertex_normals is None:
         mesh_data.compute_vertex_normals()
 
-    # 创建 Pyrender 场景并设置背景为白色
-    render_scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0])
-    mesh = pyrender.Mesh.from_trimesh(mesh_data, smooth=True)
-    mesh_node = render_scene.add(mesh)
-
-    # 设置摄像机参数
-    camera = pyrender.PerspectiveCamera(yfov=np.deg2rad(30), znear=0.0001, zfar=100000.0)
-    camera_pose = np.eye(4)
-    camera_pose[2, 3] = 4.0  # 距离模型 20 个单位
-    render_scene.add(camera, pose=camera_pose)
-
-    # 添加全局环境光
-    ambient_light = np.array([1.0, 1.0, 1.0]) * 2.0
-    render_scene.ambient_light = ambient_light
+    # 获取顶点坐标、法线和面
+    vertices = torch.tensor(mesh_data.vertices, dtype=torch.float32, device=device)
+    faces = torch.tensor(mesh_data.faces, dtype=torch.int64, device=device)
+    vertex_normals = torch.tensor(mesh_data.vertex_normals, dtype=torch.float32)
 
-    # 准备法线渲染场景
-    normals = mesh_data.vertex_normals.copy()
-
-    # 将法线映射到颜色范围 [0, 255]
-    normal_colors = ((normals + 1) / 2 * 255)
-
-    # 创建用于法线渲染的独立网格
-    normal_mesh_data = mesh_data.copy()
-    normal_mesh_data.visual.vertex_colors = np.hstack(
-        [normal_colors, np.full((normals.shape[0], 1), 255, dtype=np.uint8)]  # 添加 Alpha 通道
+    # 获取顶点颜色
+    if mesh_data.visual.vertex_colors is None:
+        # 如果没有顶点颜色，可以给定一个默认值（例如，白色）
+        vertex_colors = torch.ones_like(vertices)[None]
+    else:
+        vertex_colors = torch.tensor(mesh_data.visual.vertex_colors[:, :3], dtype=torch.float32)[None]
+    # 创建纹理并分配顶点颜色
+    textures = TexturesVertex(verts_features=vertex_colors)
+    textures.to(device)
+    # 创建Mesh对象
+    mesh = pytorch3d.structures.Meshes(verts=[vertices], faces=[faces], textures=textures)
+
+    # 设置渲染器
+    lights = AmbientLights(ambient_color=((3.0,)*3,), device=device)
+    # lights = PointLights(device=device, location=[[0.0, 0.0, 3.0]], ambient_color=[[0.5, 0.5, 0.5]], diffuse_color=[[1.0, 1.0, 1.0]])
+    raster_settings = RasterizationSettings(
+        image_size=image_size,  # 渲染图像的尺寸
+        blur_radius=0.0,  # 默认无模糊
+        faces_per_pixel=1,  # 每像素渲染一个面
+        # background_color=(1.0, 1.0, 1.0)
     )
 
-    # 创建法线渲染场景
-    normal_scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0])
-    normal_mesh = pyrender.Mesh.from_trimesh(normal_mesh_data, smooth=True)
-    normal_mesh_node = normal_scene.add(normal_mesh)
-    normal_scene.add(camera, pose=camera_pose)
-    normal_scene.ambient_light = ambient_light
-
-    # 初始化渲染器
-    r = pyrender.OffscreenRenderer(*resolution)
-
-    # 创建视频写入器
-    writer = imageio.get_writer(output_video_path, fps=fps)
-
+    # 设置旋转和渲染参数
+    frames = []
+    camera_distance = 6.5
+    elevs = 0.0
+    center = (0.0, 0.0, 0.0)
     # 渲染每一帧
-    try:
-        for frame_idx in tqdm(range(frame_count)):
-            # 计算旋转角度
-            angle = 2 * np.pi * frame_idx / frame_count
-            rotation_matrix = np.array([
-                [math.cos(angle), 0, math.sin(angle), 0],
-                [0, 1, 0, 0],
-                [-math.sin(angle), 0, math.cos(angle), 0],
-                [0, 0, 0, 1]
-            ])
-
-            # 更新模型的姿态
-            render_scene.set_pose(mesh_node, rotation_matrix)
-
-            # 渲染 RGB 图像
-            color, _ = r.render(render_scene)
-
-            # 更新法线场景的姿态
-            normal_scene.set_pose(normal_mesh_node, rotation_matrix)
-
-            # 渲染法线图像
-            normal, _ = r.render(normal_scene, flags=pyrender.RenderFlags.FLAT)
+    materials = Materials(
+            device=device,
+            diffuse_color=((0.0, 0.0, 0.0),),
+            ambient_color=((1.0, 1.0, 1.0),),
+            specular_color=((0.0, 0.0, 0.0),),
+            shininess=0.0,
+    )
         
-            # 拼接左右图像
-            combined_frame = np.concatenate((color, normal), axis=1)
-
-            # 写入视频帧
-            writer.append_data(combined_frame)
-    finally:
-        # 释放资源
-        writer.close()
-        r.delete()
-
-    print(f"Rendered video saved to {output_video_path}")
-    return output_video_path
+    rasterizer = MeshRasterizer(raster_settings=raster_settings)
+    for i in tqdm(range(num_frames)):
+        azims = 360.0 * i / num_frames
+        R, T = look_at_view_transform(
+            dist=camera_distance,
+            elev=elevs,
+            azim=azims,
+            at=(center,),
+            degrees=True
+        )
+
+
+        # 手动设置相机的旋转矩阵
+        cameras = PerspectiveCameras(device=device, R=R, T=T, focal_length=5.0)
+        cameras.znear = 0.0001
+        cameras.zfar = 10000000.0
+        shader=MultiOutputShader(
+                device=device,
+                cameras=cameras,
+                lights=lights,
+                materials=materials,
+                choices=["rgb", "mask", "normal"]
+            )
+
+        renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
+        # 渲染RGB图像和Normal图像
+        render_result = renderer(mesh, cameras=cameras)
+        rgb_image = render_result["rgb"] * render_result["mask"] + (1 - render_result["mask"]) * torch.ones_like(render_result["rgb"]) * 255.0
+        normal_map = render_result["normal"]
+
+        # 提取RGB和Normal map
+        rgb = rgb_image[0, ..., :3].cpu().numpy()  # RGB图像
+        normal_map = torch.nn.functional.normalize(normal_map, dim=-1)  # Normal map
+        normal_map = (normal_map + 1) / 2
+        normal_map = normal_map * render_result["mask"] + (1 - render_result["mask"]) * torch.ones_like(render_result["normal"])
+        normal = normal_map[0, ..., :3].cpu().numpy()  # Normal map
+        rgb = np.clip(rgb, 0, 255).astype(np.uint8)
+        normal = np.clip(normal*255, 0, 255).astype(np.uint8)
+        # 将RGB和Normal map合并为一张图，左边RGB，右边Normal map
+        combined_image = np.concatenate((rgb, normal), axis=1)
+
+        # 将合并后的图像加入到帧列表
+        frames.append(combined_image)
+
+    # 使用imageio保存视频
+    imageio.mimsave(output_video_path, frames, fps=fps)
+
+    print(f"Video saved to {output_video_path}")
 
 if __name__ == '__main__':
     # 示例调用
-    input_obj_path = "output/gradio_cache/text_3D/_超级赛亚人_10/rgb_projected.obj"
+    input_obj_path = "/hpc2hdd/home/jlin695/code/github/Kiss3DGen/outputs/a_owl_wearing_a_hat/ISOMER/rgb_projected.obj"
     output_video_path = "output.mp4"
     render_video_from_obj(input_obj_path, output_video_path)