pytorch3D如何输出渲染图像的深度图

在这个包里修改，获得能够输出深度的渲染器
/home/fuxinwen/miniconda3/envs/pytorch3d/lib/python3.9/site-packages/pytorch3d/renderer/mesh/shader.py

class SoftPhongShaderWithDepth(ShaderBase):
    """
    每像素光照 - 该光照模型通过每个像素的插值坐标和法线应用。
    混合函数返回每个像素所有面的软聚合颜色。

    使用默认值时，可以按以下方式初始化着色器：

    .. code-block::

        shader = SoftPhongShader(device=torch.device("cuda:0"))
    """

    def forward(self, fragments: Fragments, meshes: Meshes, **kwargs) -> (torch.Tensor, torch.Tensor):
        # 从超类获取相机对象
        cameras = super()._get_cameras(**kwargs)
        
        # 从网格中采样纹理
        texels = meshes.sample_textures(fragments)
        
        # 获取光源、材料和混合参数
        lights = kwargs.get("lights", self.lights)
        materials = kwargs.get("materials", self.materials)
        blend_params = kwargs.get("blend_params", self.blend_params)
        
        # 进行 Phong 着色计算
        colors = phong_shading(
            meshes=meshes,
            fragments=fragments,
            texels=texels,
            lights=lights,
            cameras=cameras,
            materials=materials,
        )
        
        # 获取相机的近裁剪面和远裁剪面
        znear = kwargs.get("znear", getattr(cameras, "znear", 1.0))
        zfar = kwargs.get("zfar", getattr(cameras, "zfar", 100.0))
        
        # 进行软最大 RGB 混合
        images = softmax_rgb_blend(
            colors, fragments, blend_params, znear=znear, zfar=zfar
        )

        # 输出深度信息
        depths = fragments.zbuf  # 假设 fragments 对象中有深度缓冲区 zbuf
        
        return images, depths

class HardPhongShaderWithDepth(ShaderBase):
    """
    Per pixel lighting - the lighting model is applied using the interpolated
    coordinates and normals for each pixel. The blending function hard assigns
    the color of the closest face for each pixel.

    To use the default values, simply initialize the shader with the desired
    device e.g.

    .. code-block::

        shader = HardPhongShader(device=torch.device("cuda:0"))
    """

    def forward(self, fragments: Fragments, meshes: Meshes, **kwargs) -> torch.Tensor:
        cameras = super()._get_cameras(**kwargs)
        texels = meshes.sample_textures(fragments)
        lights = kwargs.get("lights", self.lights)
        materials = kwargs.get("materials", self.materials)
        blend_params = kwargs.get("blend_params", self.blend_params)
        colors = phong_shading(
            meshes=meshes,
            fragments=fragments,
            texels=texels,
            lights=lights,
            cameras=cameras,
            materials=materials,
        )
        images = hard_rgb_blend(colors, fragments, blend_params)
        # 输出深度信息
        depths = fragments.zbuf  # 假设 fragments 对象中有深度缓冲区 zbuf
        
        return images, depths

然后将渲染器注册到下面两个__init__.py文件中，就可以在包中调用上面两个渲染函数了
/home/fuxinwen/miniconda3/envs/pytorch3d/lib/python3.9/site-packages/pytorch3d/renderer/init.py

/home/fuxinwen/miniconda3/envs/pytorch3d/lib/python3.9/site-packages/pytorch3d/renderer/mesh/init.py