Spaces:
mashroo
/
Running on Zero

YoussefAnso commited on
Commit
30b6055
·
1 Parent(s): 76eeb7d

Refactor generate3d function in inference.py to improve mesh export process. Updated OBJ and texture file checks, streamlined GLB conversion with embedded texture handling, and enhanced error handling for missing files. This change ensures proper export of textured GLB files.

Browse files
Files changed (1) hide show
  1. inference.py +13 -22
inference.py CHANGED
@@ -4,7 +4,6 @@ import time
4
  import tempfile
5
  import os
6
  from PIL import Image
7
- import zipfile
8
  import trimesh
9
 
10
  from util.utils import get_tri
@@ -63,7 +62,6 @@ def generate3d(model, rgb, ccm, device):
63
  data_config['verts'] = verts[0]
64
  data_config['faces'] = faces
65
 
66
- # Clean the mesh
67
  from kiui.mesh_utils import clean_mesh
68
  verts, faces = clean_mesh(
69
  data_config['verts'].squeeze().cpu().numpy().astype(np.float32),
@@ -73,10 +71,10 @@ def generate3d(model, rgb, ccm, device):
73
  data_config['verts'] = torch.from_numpy(verts).cuda().contiguous()
74
  data_config['faces'] = torch.from_numpy(faces).cuda().contiguous()
75
 
76
- # Export mesh with UV
77
  mesh_path_obj = tempfile.NamedTemporaryFile(suffix="", delete=False).name
78
  model.export_mesh_wt_uv(
79
- None, # GL context not needed for CPU export
80
  data_config,
81
  mesh_path_obj,
82
  "",
@@ -85,27 +83,20 @@ def generate3d(model, rgb, ccm, device):
85
  tri_fea_2=triplane_feature2
86
  )
87
 
88
- # Check if texture exists
 
89
  texture_path = mesh_path_obj + ".png"
90
- if not os.path.exists(texture_path):
91
- raise RuntimeError("Texture image not created, cannot export textured GLB.")
92
 
93
- # Load the .obj file and apply the texture manually
94
- scene_or_mesh = trimesh.load(mesh_path_obj + ".obj", force='scene')
95
 
96
- texture_image = Image.open(texture_path)
97
- texture_image = np.array(texture_image)
98
 
99
- mesh_path_glb = mesh_path_obj + "_textured.glb"
 
100
 
101
- if isinstance(scene_or_mesh, trimesh.Scene):
102
- for _, geometry in scene_or_mesh.geometry.items():
103
- material = trimesh.visual.texture.SimpleMaterial(image=texture_image)
104
- geometry.visual = trimesh.visual.texture.TextureVisuals(image=texture_image, material=material)
105
- scene_or_mesh.export(mesh_path_glb)
106
- else:
107
- material = trimesh.visual.texture.SimpleMaterial(image=texture_image)
108
- scene_or_mesh.visual = trimesh.visual.texture.TextureVisuals(image=texture_image, material=material)
109
- scene_or_mesh.export(mesh_path_glb)
110
 
111
- return mesh_path_glb
 
4
  import tempfile
5
  import os
6
  from PIL import Image
 
7
  import trimesh
8
 
9
  from util.utils import get_tri
 
62
  data_config['verts'] = verts[0]
63
  data_config['faces'] = faces
64
 
 
65
  from kiui.mesh_utils import clean_mesh
66
  verts, faces = clean_mesh(
67
  data_config['verts'].squeeze().cpu().numpy().astype(np.float32),
 
71
  data_config['verts'] = torch.from_numpy(verts).cuda().contiguous()
72
  data_config['faces'] = torch.from_numpy(faces).cuda().contiguous()
73
 
74
+ # Export to .obj with UV and texture image
75
  mesh_path_obj = tempfile.NamedTemporaryFile(suffix="", delete=False).name
76
  model.export_mesh_wt_uv(
77
+ None, # GL context skipped (no dr.RasterizeGLContext)
78
  data_config,
79
  mesh_path_obj,
80
  "",
 
83
  tri_fea_2=triplane_feature2
84
  )
85
 
86
+ # Convert to GLB with embedded texture
87
+ obj_path = mesh_path_obj + ".obj"
88
  texture_path = mesh_path_obj + ".png"
 
 
89
 
90
+ if not os.path.exists(obj_path) or not os.path.exists(texture_path):
91
+ raise RuntimeError("OBJ or texture file missing — cannot convert to GLB.")
92
 
93
+ mesh_scene = trimesh.load(obj_path, force='scene')
94
+ texture_image = np.array(Image.open(texture_path))
95
 
96
+ for name, geom in mesh_scene.geometry.items():
97
+ geom.visual = trimesh.visual.texture.TextureVisuals(image=texture_image)
98
 
99
+ glb_path = mesh_path_obj + "_final.glb"
100
+ mesh_scene.export(glb_path)
 
 
 
 
 
 
 
101
 
102
+ return glb_path