obj and png outputs work

gradio_app.py

@@ -1,5 +1,6 @@
 import os
 import random
+from re import I
 import tempfile
 import time
 import zipfile
@@ -65,6 +66,10 @@ example_files = [
 
 
 def create_zip_file(glb_file, pc_file, illumination_file):
+    print("zip disabled for perf")
+
+    return "zip disabled for perf"
+
     if not all([glb_file, pc_file, illumination_file]):
         return None
 
@@ -206,22 +211,47 @@ def run_model(
 
     # Create new tmp file
     temp_dir = tempfile.mkdtemp()
-    tmp_file = os.path.join(temp_dir, "mesh.glb")
+    tmp_file = os.path.join(temp_dir, "mesh.obj")
 
-    trimesh_mesh.export(tmp_file, file_type="glb", include_normals=True)
+    trimesh_mesh.export(tmp_file, file_type="obj", include_normals=True)
     generated_files.append(tmp_file)
 
-    tmp_file_pc = os.path.join(temp_dir, "points.ply")
-    trimesh_pc.export(tmp_file_pc)
-    generated_files.append(tmp_file_pc)
-
-    tmp_file_illumination = os.path.join(temp_dir, "illumination.hdr")
-    cv2.imwrite(tmp_file_illumination, illumination_map)
-    generated_files.append(tmp_file_illumination)
+    # tmp_file_pc = os.path.join(temp_dir, "points.ply")
+    # trimesh_pc.export(tmp_file_pc)
+    # generated_files.append(tmp_file_pc)
+
+    # tmp_file_illumination = os.path.join(temp_dir, "illumination.hdr")
+    # cv2.imwrite(tmp_file_illumination, illumination_map)
+    # generated_files.append(tmp_file_illumination)
+
+    # Extract textures
+    texture_full_path = None
+    if trimesh_mesh.visual is not None:
+        material = trimesh_mesh.visual.material
+        if hasattr(material, 'baseColorTexture'):
+            texture_path = 'texture.png'
+            texture_full_path = os.path.join(temp_dir, texture_path)
+            
+            # Convert the texture data to an image and save
+            texture_data = material.baseColorTexture
+            if isinstance(texture_data, np.ndarray):
+                texture_img = Image.fromarray(texture_data)
+                print("Saving texture to", texture_full_path)
+                texture_img.save(texture_full_path)
+                generated_files.append(texture_full_path)
+            elif isinstance(texture_data, Image.Image):
+                print("Saving texture to", texture_full_path)
+                texture_data.save(texture_full_path)
+                generated_files.append(texture_full_path)
+            else:
+                print("Texture data is not a numpy array, but instead: " + str(type(texture_data)))
+        else:
+            print("Material has no baseColorTexture, but only:")
+            print(vars(material))
 
     print("Generation took:", time.time() - start, "s")
 
-    return tmp_file, tmp_file_pc, tmp_file_illumination, trimesh_pc
+    return tmp_file, None, None, trimesh_pc, texture_full_path
 
 
 def create_batch(input_image: Image) -> dict[str, Any]:
@@ -300,7 +330,7 @@ def process_model_run(
         f"Final vertex count: {final_vertex_count} with type {vertex_count_type} and vertex count {vertex_count}"
     )
 
-    glb_file, pc_file, illumination_file, pc_plot = run_model(
+    glb_file, pc_file, illumination_file, pc_plot, texture_file = run_model(
         background_state,
         guidance_scale,
         random_seed,
@@ -323,7 +353,7 @@ def process_model_run(
             ]
         )
 
-    return glb_file, pc_file, illumination_file, point_list
+    return glb_file, pc_file, illumination_file, point_list, texture_file
 
 
 def regenerate_run(
@@ -336,7 +366,7 @@ def regenerate_run(
     vertex_count,
     texture_resolution,
 ):
-    glb_file, pc_file, illumination_file, point_list = process_model_run(
+    glb_file, pc_file, illumination_file, point_list, texture_file = process_model_run(
         background_state,
         guidance_scale,
         random_seed,
@@ -394,7 +424,7 @@ def run_button(
         else:
             pc_cond = None
 
-        glb_file, pc_file, illumination_file, pc_list = process_model_run(
+        glb_file, pc_file, illumination_file, pc_list, texture_file = process_model_run(
             background_state,
             guidance_scale,
             random_seed,
@@ -405,7 +435,8 @@ def run_button(
             texture_resolution,
         )
 
-        zip_file = create_zip_file(glb_file, pc_file, illumination_file)
+        # Disabled to improve performance
+        # zip_file = create_zip_file(glb_file, pc_file, illumination_file)
 
         if torch.cuda.is_available():
             print("Peak Memory:", torch.cuda.max_memory_allocated() / 1024 / 1024, "MB")
@@ -426,7 +457,8 @@ def run_button(
             gr.update(value=pc_list),  # point_cloud_editor
             gr.update(value=pc_file),  # pc_download
             gr.update(visible=False),  # regenerate_btn
-            gr.update(value=zip_file, visible=True),  # download_all_btn
+            gr.update(value=None, visible=True),  # download_all_btn
+            gr.update(value=texture_file, visible=True),  # download_all_btn
         )
 
     elif run_btn == "Remove Background":
@@ -725,6 +757,14 @@ with gr.Blocks() as demo:
                 contrast=1.0,
                 scale=1.0,
             )
+
+            output_texture = gr.File(
+                label="Download Texture",
+                file_types=[".png"],
+                file_count="single",
+                visible=False,
+            )
+
             with gr.Column(visible=False, scale=1.0) as hdr_row:
                 gr.Markdown(
                     """## HDR Environment Map
@@ -875,6 +915,7 @@ with gr.Blocks() as demo:
             pc_download,
             regenerate_btn,
             download_all_btn,
+            output_texture,
         ],
     )
 

run.py

@@ -10,6 +10,7 @@ from transparent_background import Remover
 from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
 from spar3d.system import SPAR3D
 from spar3d.utils import foreground_crop, get_device, remove_background
+import trimesh
 
 
 def check_positive(value):
@@ -178,13 +179,22 @@ if __name__ == "__main__":
             )
 
         if len(image) == 1:
-            out_mesh_path = os.path.join(output_dir, str(i), "mesh.glb")
+            assert(isinstance(mesh, trimesh.Trimesh))
+            out_mesh_path = os.path.join(output_dir, str(i), "mesh.obj")
             mesh.export(out_mesh_path, include_normals=True)
             out_points_path = os.path.join(output_dir, str(i), "points.ply")
             glob_dict["point_clouds"][0].export(out_points_path)
+
+            # Extract textures
+            for j, material in enumerate(mesh.visual.material):
+                if hasattr(material, 'image'):
+                    texture_path = f'texture_{j}.png'
+                    with open(os.path.join(output_dir, str(i),  texture_path), 'wb') as f:
+                        f.write(material.image.tobytes())
+
         else:
             for j in range(len(mesh)):
-                out_mesh_path = os.path.join(output_dir, str(i + j), "mesh.glb")
+                out_mesh_path = os.path.join(output_dir, str(i + j), "mesh.obj")
                 mesh[j].export(out_mesh_path, include_normals=True)
                 out_points_path = os.path.join(output_dir, str(i + j), "points.ply")
                 glob_dict["point_clouds"][j].export(out_points_path)