Enhance texture baking logic in generate3d function of inference.py by integrating OpenCV for color map processing. Updated albedo generation to support photorealistic textures and ensure proper format handling, improving overall texture quality in mesh exports.

inference.py

@@ -5,6 +5,7 @@ import tempfile
 import zipfile
 import nvdiffrast.torch as dr
 import xatlas
+import cv2
 
 from util.utils import get_tri
 from mesh import Mesh
@@ -66,10 +67,19 @@ def generate3d(model, rgb, ccm, device):
     mesh_f = data_config['faces'].cpu().numpy()
     vmapping, ft, vt = xatlas.parametrize(mesh_v, mesh_f)
 
-    # Bake texture (simulate what export_mesh_wt_uv does, but CPU-only)
-    # Here, we just fill with white for demo; replace with your actual texture baking logic
-    tex_res = (1024, 1024)
-    albedo = np.ones((tex_res[0], tex_res[1], 3), dtype=np.float32)  # TODO: bake your texture here
+    # Load your color map (photorealistic texture)
+    # Assume you save the color map as a PNG somewhere, or pass it as a numpy array
+    # If rgb is your color map, convert it to the right format:
+    if rgb.max() > 1.0:
+        color_map = rgb.astype(np.uint8)
+    else:
+        color_map = (rgb * 255).astype(np.uint8)
+    # If rgb is already in HWC, skip transpose; otherwise, transpose as needed
+    if color_map.shape[-1] != 3:
+        color_map = np.transpose(color_map, (1, 2, 0))
+    albedo = cv2.cvtColor(color_map, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+
+    tex_res = (albedo.shape[0], albedo.shape[1])
 
     # Create Mesh and export .glb
     mesh = Mesh(