diff --git a/trellis/renderers/mesh_renderer.py b/trellis/renderers/mesh_renderer.py index b504fa4d..085f1565 100644 --- a/trellis/renderers/mesh_renderer.py +++ b/trellis/renderers/mesh_renderer.py @@ -57,7 +57,7 @@ def render( mesh : MeshExtractResult, extrinsics: torch.Tensor, intrinsics: torch.Tensor, - return_types = ["mask", "normal", "depth"] + return_types = ["mask", "normal", "depth", "color"] ) -> edict: """ Render the mesh. diff --git a/trellis/utils/postprocessing_utils.py b/trellis/utils/postprocessing_utils.py index 0a8d9fb7..3840b0f8 100644 --- a/trellis/utils/postprocessing_utils.py +++ b/trellis/utils/postprocessing_utils.py @@ -397,7 +397,7 @@ def tv_loss(texture): def to_glb( - app_rep: Union[Strivec, Gaussian], + app_rep: Union[Strivec, Gaussian, MeshExtractResult], mesh: MeshExtractResult, simplify: float = 0.95, fill_holes: bool = True, @@ -410,7 +410,7 @@ def to_glb( Convert a generated asset to a glb file. Args: - app_rep (Union[Strivec, Gaussian]): Appearance representation. + app_rep (Union[Strivec, Gaussian, MeshExtractResult]): Appearance representation. mesh (MeshExtractResult): Extracted mesh. simplify (float): Ratio of faces to remove in simplification. fill_holes (bool): Whether to fill holes in the mesh. diff --git a/trellis/utils/render_utils.py b/trellis/utils/render_utils.py index 8187c84f..46b0a2fd 100644 --- a/trellis/utils/render_utils.py +++ b/trellis/utils/render_utils.py @@ -82,7 +82,9 @@ def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=N rets['depth'].append(None) else: res = renderer.render(sample, extr, intr) + if 'color' not in rets: rets['color'] = [] if 'normal' not in rets: rets['normal'] = [] + rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) return rets