debug

inference.py

@@ -1,13 +1,15 @@
 import numpy as np
 import torch
 import time
-import nvdiffrast.torch as dr
-from util.utils import get_tri
 import tempfile
-from mesh import Mesh
+import os
+from PIL import Image
 import zipfile
+import trimesh
+
+from util.utils import get_tri
+from mesh import Mesh
 from util.renderer import Renderer
-import trimesh  # Needed for glb export
 
 def generate3d(model, rgb, ccm, device):
     model.renderer = Renderer(
@@ -17,17 +19,18 @@ def generate3d(model, rgb, ccm, device):
         geo_type=model.geo_type
     )
 
-    # RGB and coordinate conversion
     color_tri = torch.from_numpy(rgb) / 255
     xyz_tri = torch.from_numpy(ccm[:, :, (2, 1, 0)]) / 255
+
     color = color_tri.permute(2, 0, 1)
     xyz = xyz_tri.permute(2, 0, 1)
 
-    def get_imgs(color):
-        color_list = [color[:, :, 256 * 5:256 * (1 + 5)]]
-        for i in range(0, 5):
-            color_list.append(color[:, :, 256 * i:256 * (1 + i)])
-        return torch.stack(color_list, dim=0)
+    def get_imgs(img_tensor):
+        images = []
+        images.append(img_tensor[:, :, 256*5:256*(1+5)])
+        for i in range(5):
+            images.append(img_tensor[:, :, 256*i:256*(i+1)])
+        return torch.stack(images, dim=0)  # [6, C, H, W]
 
     triplane_color = get_imgs(color).permute(0, 2, 3, 1).unsqueeze(0).to(device)
     color = get_imgs(color)
@@ -47,19 +50,20 @@ def generate3d(model, rgb, ccm, device):
         with torch.no_grad():
             triplane_feature2 = model.unet2(triplane, tnew)
     else:
-        triplane_feature2 = model.unet2(triplane)
+        with torch.no_grad():
+            triplane_feature2 = model.unet2(triplane)
 
-    with torch.no_grad():
-        data_config = {
-            'resolution': [1024, 1024],
-            'triview_color': triplane_color.to(device),
-        }
+    data_config = {
+        'resolution': [1024, 1024],
+        "triview_color": triplane_color.to(device),
+    }
 
+    with torch.no_grad():
         verts, faces = model.decode(data_config, triplane_feature2)
         data_config['verts'] = verts[0]
         data_config['faces'] = faces
 
-    # Optional mesh cleanup (reduce remesh for speed)
+    # Clean the mesh
     from kiui.mesh_utils import clean_mesh
     verts, faces = clean_mesh(
         data_config['verts'].squeeze().cpu().numpy().astype(np.float32),
@@ -69,23 +73,39 @@ def generate3d(model, rgb, ccm, device):
     data_config['verts'] = torch.from_numpy(verts).cuda().contiguous()
     data_config['faces'] = torch.from_numpy(faces).cuda().contiguous()
 
-    # Rasterization context
-    glctx = dr.RasterizeGLContext()
-
-    # Temporary output path
+    # Export mesh with UV
     mesh_path_obj = tempfile.NamedTemporaryFile(suffix="", delete=False).name
+    model.export_mesh_wt_uv(
+        None,  # GL context not needed for CPU export
+        data_config,
+        mesh_path_obj,
+        "",
+        device,
+        res=(1024, 1024),
+        tri_fea_2=triplane_feature2
+    )
 
-    # Export OBJ with UV and PNG
-    with torch.no_grad():
-        model.export_mesh_wt_uv(
-            glctx, data_config, mesh_path_obj, "", device,
-            res=(512, 512), tri_fea_2=triplane_feature2
-        )
+    # Check if texture exists
+    texture_path = mesh_path_obj + ".png"
+    if not os.path.exists(texture_path):
+        raise RuntimeError("Texture image not created, cannot export textured GLB.")
 
-    # Convert to .glb using trimesh
-    mesh = trimesh.load(mesh_path_obj + ".obj", force='mesh')
-    mesh_path_glb = mesh_path_obj + ".glb"
-    mesh.export(mesh_path_glb, file_type='glb')
+    # Load the .obj file and apply the texture manually
+    scene_or_mesh = trimesh.load(mesh_path_obj + ".obj", force='scene')
+
+    texture_image = Image.open(texture_path)
+    texture_image = np.array(texture_image)
+
+    mesh_path_glb = mesh_path_obj + "_textured.glb"
+
+    if isinstance(scene_or_mesh, trimesh.Scene):
+        for _, geometry in scene_or_mesh.geometry.items():
+            material = trimesh.visual.texture.SimpleMaterial(image=texture_image)
+            geometry.visual = trimesh.visual.texture.TextureVisuals(image=texture_image, material=material)
+        scene_or_mesh.export(mesh_path_glb)
+    else:
+        material = trimesh.visual.texture.SimpleMaterial(image=texture_image)
+        scene_or_mesh.visual = trimesh.visual.texture.TextureVisuals(image=texture_image, material=material)
+        scene_or_mesh.export(mesh_path_glb)
 
-    print(f"✅ Exported GLB with UV texture: {mesh_path_glb}")
     return mesh_path_glb