Fix depth estimation

app.py

@@ -100,7 +100,7 @@ from utils.version_info import (
     #release_torch_resources,
     #get_torch_info
 )
-
+from utils.depth_estimation import (get_depth_map_from_state)
 
 input_image_palette = []
 current_prerendered_image = gr.State("./images/images/Beeuty-1.png")
@@ -722,7 +722,8 @@ def generate_3d_asset(depth_image_source, randomize_seed, seed, input_image, out
     # Validate the mesh
     mesh = outputs['mesh'][0]
     # Depending on the mesh format (it might be a dict or an object)
-    if isinstance(mesh, dict):
+    meshisdict = isinstance(mesh, dict)
+    if meshisdict:
         vertices = mesh['vertices']
         faces = mesh['faces']
     else:
@@ -738,20 +739,29 @@ def generate_3d_asset(depth_image_source, randomize_seed, seed, input_image, out
     if not vertices.is_cuda or not faces.is_cuda:
         raise ValueError("Mesh data must be on GPU")
     if vertices.dtype != torch.float32 or faces.dtype != torch.int32:
-        raise ValueError("Mesh vertices must be float32, faces must be int32")
+        if meshisdict:
+            mesh['faces'] = faces.to(torch.int32)
+            mesh['vertices'] = vertices.to(torch.float32)
+        else:
+            mesh.faces = faces.to(torch.int32)
+            mesh.vertices = vertices.to(torch.float32)
+        #raise ValueError("Mesh vertices must be float32, faces must be int32")
 
     # Save the video to a temporary file
     user_dir = os.path.join(constants.TMPDIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
 
     video = render_utils.render_video(outputs['gaussian'][0], resolution=576, num_frames=60, r=1)['color']
-    snapshot_results = render_utils.render_snapshot(outputs['gaussian'][0], resolution=576)
-    depth_snapshot = snapshot_results['depth'][0]
     video_geo = render_utils.render_video(outputs['mesh'][0], resolution=576, num_frames=60, r=1)['normal']
     video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
     video_path = os.path.join(user_dir, f'{output_name}.mp4')
-    imageio.mimsave(video_path, video, fps=15)
+    imageio.mimsave(video_path, video, fps=10)
+
+    snapshot_results = render_utils.render_snapshot(outputs['mesh'][0], resolution=768)['depth']
+    depth_snapshot = Image.fromarray(snapshot_results[0])
+
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], output_name)
+    #depth_snapshot = get_depth_map_from_state(state, image_raw.size[0], image_raw.size[1])
     torch.cuda.empty_cache()
     return [state, video_path, depth_snapshot]
 
@@ -1106,7 +1116,7 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
         with gr.Row():
             with gr.Column():
                 # Use standard seed settings only
-                seed_3d = gr.Slider(0, constants.MAX_SEED, label="Seed (3D Generation)", value=0, step=1)
+                seed_3d = gr.Slider(0, constants.MAX_SEED, label="Seed (3D Generation)", value=0, step=1, randomize=True)
                 randomize_seed_3d = gr.Checkbox(label="Randomize Seed (3D Generation)", value=True)
             with gr.Column():
                 depth_image_source = gr.Radio(

trellis/renderers/gaussian_render.py

@@ -11,7 +11,6 @@
 
 import torch
 import math
-from easydict import EasyDict as edict
 import numpy as np
 from ..representations.gaussian import Gaussian
 from .sh_utils import eval_sh

trellis/utils/render_utils.py

@@ -67,6 +67,53 @@ def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=N
     else:
         raise ValueError(f'Unsupported sample type: {type(sample)}')
     
+    rets = {}
+    for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), desc='Rendering', disable=not verbose):
+        if not isinstance(sample, MeshExtractResult):
+            res = renderer.render(sample, extr, intr, colors_overwrite=colors_overwrite)
+            if 'color' not in rets: rets['color'] = []
+            # if 'depth' not in rets: rets['depth'] = []
+            rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
+            # if 'percent_depth' in res:
+            #     rets['depth'].append(res['percent_depth'].detach().cpu().numpy())
+            # elif 'depth' in res:
+            #     rets['depth'].append(res['depth'].detach().cpu().numpy())
+            # else:
+            #     rets['depth'].append(None)
+        else:
+            res = renderer.render(sample, extr, intr)
+            if 'normal' not in rets: rets['normal'] = []
+            rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
+
+    return rets
+
+def render_frames_depth(sample, extrinsics, intrinsics, options={}, colors_overwrite=None, verbose=True, **kwargs):
+    if isinstance(sample, Octree):
+        renderer = OctreeRenderer()
+        renderer.rendering_options.resolution = options.get('resolution', 512)
+        renderer.rendering_options.near = options.get('near', 0.8)
+        renderer.rendering_options.far = options.get('far', 1.6)
+        renderer.rendering_options.bg_color = options.get('bg_color', (0, 0, 0))
+        renderer.rendering_options.ssaa = options.get('ssaa', 4)
+        renderer.pipe.primitive = sample.primitive
+    elif isinstance(sample, Gaussian):
+        renderer = GaussianRenderer()
+        renderer.rendering_options.resolution = options.get('resolution', 512)
+        renderer.rendering_options.near = options.get('near', 0.8)
+        renderer.rendering_options.far = options.get('far', 1.6)
+        renderer.rendering_options.bg_color = options.get('bg_color', (0, 0, 0))
+        renderer.rendering_options.ssaa = options.get('ssaa', 1)
+        renderer.pipe.kernel_size = kwargs.get('kernel_size', 0.1)
+        renderer.pipe.use_mip_gaussian = True
+    elif isinstance(sample, MeshExtractResult):
+        renderer = MeshRenderer()
+        renderer.rendering_options.resolution = options.get('resolution', 512)
+        renderer.rendering_options.near = options.get('near', 1)
+        renderer.rendering_options.far = options.get('far', 100)
+        renderer.rendering_options.ssaa = options.get('ssaa', 4)
+    else:
+        raise ValueError(f'Unsupported sample type: {type(sample)}')
+    
     rets = {}
     for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), desc='Rendering', disable=not verbose):
         if not isinstance(sample, MeshExtractResult):
@@ -82,12 +129,14 @@ def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=N
                 rets['depth'].append(None)
         else:
             res = renderer.render(sample, extr, intr)
-            if 'normal' not in rets: rets['normal'] = []
-            rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
+            if 'depth' not in rets: rets['depth'] = []
+            if 'depth' in res:
+                rets['depth'].append(np.clip(res['depth'].detach().cpu().numpy(), 0, 255).astype(np.uint8))
+            else:
+                rets['depth'].append(None)
     return rets
 
-
-def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, **kwargs):
+def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=60, **kwargs):
     yaws = torch.linspace(0, 2 * 3.1415, num_frames)
     pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
     yaws = yaws.tolist()
@@ -114,3 +163,11 @@ def render_snapshot(samples, resolution=512, bg_color=(0, 0, 0), offset=(-16 / 1
     pitch = [offset[1] for _ in range(4)]
     extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, r, fov)
     return render_frames(samples, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
+
+def render_snapshot_depth(samples, resolution=512, bg_color=(0, 0, 0), offset=(0, np.pi/2), r=1,fov=80, **kwargs):
+    yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+    yaw_offset = offset[0]
+    yaw = [y + yaw_offset for y in yaw]
+    pitch = [offset[1] for _ in range(4)]
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, r, fov)
+    return render_frames_depth(samples, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)

utils/depth_estimation.py

@@ -12,6 +12,8 @@ from utils.image_utils import (
     resize_image_with_aspect_ratio
 )
 from utils.constants import TMPDIR
+from easydict import EasyDict as edict
+
 
 # Load models once during module import
 image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
@@ -258,10 +260,10 @@ def depth_process_image(image_path, resized_width=800, z_scale=208):
         torch.cuda.ipc_collect()
     return [img, gltf_path, gltf_path]
 
-def get_depth_map_from_state(state):
+def get_depth_map_from_state(state, image_height=1024, image_width=1024):
     from diff_gaussian_rasterization import GaussianRasterizer, GaussianRasterizationSettings
 
-    settings = GaussianRasterizationSettings(image_height=1024, image_width=1024, bg_color=(1.0, 1.0, 1.0), max_num_points_per_tile=100, tile_size=32, filter_mode="linear", precompute_cov3D=True, precompute_cov2D=True, precompute_colors=True)
+    settings = GaussianRasterizationSettings(image_height=image_height, image_width=image_width, kernel_size=0.01,bg=(0.0, 0.0, 0.0))
     rasterizer = GaussianRasterizer(settings)
     # Assume state has necessary data like means3D, scales, etc.
     rendered_image, rendered_depth, _, _, _, _ = rasterizer(means3D=state["means3D"], means2D=state["means2D"], shs=state["shs"], colors_precomp=state["colors_precomp"], opacities=state["opacities"], scales=state["scales"], rotations=state["rotations"], cov3D_precomp=state["cov3D_precomp"])