Refactor generate3d function in inference.py to improve mesh export process. Introduced temporary file handling for OBJ, MTL, and PNG outputs, ensuring proper texture mapping. Enhanced code readability by restructuring renderer initialization and file writing logic.

inference.py

@@ -1,16 +1,20 @@
 import numpy as np
 import torch
 import time
-import nvdiffrast.torch as dr
 from util.utils import get_tri
 import tempfile
 from util.renderer import Renderer
 import os
+from PIL import Image
 
-def generate3d(model, rgb, ccm, device):
 
-    model.renderer = Renderer(tet_grid_size=model.tet_grid_size, camera_angle_num=model.camera_angle_num,
-                                scale=model.input.scale, geo_type=model.geo_type)
+def generate3d(model, rgb, ccm, device):
+    model.renderer = Renderer(
+        tet_grid_size=model.tet_grid_size,
+        camera_angle_num=model.camera_angle_num,
+        scale=model.input.scale,
+        geo_type=model.geo_type
+    )
 
     color_tri = torch.from_numpy(rgb) / 255
     xyz_tri = torch.from_numpy(ccm[:, :, (2, 1, 0)]) / 255
@@ -64,11 +68,27 @@ def generate3d(model, rgb, ccm, device):
     data_config['verts'] = torch.from_numpy(verts).to(device).contiguous()
     data_config['faces'] = torch.from_numpy(faces).to(device).contiguous()
 
-    with torch.no_grad():
-        mesh_path_base = tempfile.NamedTemporaryFile(suffix="", delete=False).name
+    # === Export OBJ/MTL/PNG ===
+    obj_path = tempfile.NamedTemporaryFile(suffix=".obj", delete=False).name
+    base_path = obj_path[:-4]  # remove .obj
+
+    texture_path = base_path + ".png"
+    mtl_path = base_path + ".mtl"
+
+    model.export_mesh_geometry(data_config, obj_path)  # writes .obj with UVs
+    model.export_texture_image(data_config, texture_path)  # saves PNG
+
+    # Write MTL file manually
+    with open(mtl_path, "w") as f:
+        f.write("newmtl material0\n")
+        f.write("Kd 1.000000 1.000000 1.000000\n")
+        f.write(f"map_Kd {os.path.basename(texture_path)}\n")
 
-        # Export mesh with UV, texture, and MTL
-        ctx = dr.RasterizeCudaContext(device=device)
-        model.export_mesh_wt_uv(ctx, data_config, mesh_path_base, ind=0, device=device, res=(1024, 1024), tri_fea_2=triplane_feature2)
+    # Append .mtl reference to OBJ file
+    with open(obj_path, "r") as original:
+        lines = original.readlines()
+    with open(obj_path, "w") as modified:
+        modified.write(f"mtllib {os.path.basename(mtl_path)}\n")
+        modified.writelines(lines)
 
-    return mesh_path_base + ".obj"
+    return obj_path