use trimesh

app.py

@@ -4,14 +4,18 @@ from pathlib import Path
 import uuid
 import tempfile
 from typing import Union
-import spaces
 import atexit
+import spaces
 from concurrent.futures import ThreadPoolExecutor
 
 import gradio as gr
 import cv2
 import torch
 import numpy as np
+import click
+import trimesh
+import trimesh.visual
+from PIL import Image
 
 from moge.model import MoGeModel
 from moge.utils.vis import colorize_depth
@@ -33,47 +37,68 @@ def delete_later(path: Union[str, os.PathLike], delay: int = 300):
     thread_pool_executor.submit(_wait_and_delete)
     atexit.register(_delete)
 
+
 @spaces.GPU
-def run(image: np.ndarray, remove_edge: bool = True):
+def run_with_gpu(image: np.ndarray):
+    image_tensor = torch.tensor(image, dtype=torch.float32, device=torch.device('cuda')).permute(2, 0, 1) / 255
+    output = model.infer(image_tensor, resolution_level=9, apply_mask=True)
+    output = {k: v.cpu().numpy() for k, v in output.items()}
+    return output
+
+
+def run(image: np.ndarray, remove_edge: bool = True, max_size: int = 800):
     run_id = str(uuid.uuid4())
 
     larger_size = max(image.shape[:2])
-    if larger_size > 1024:
-        scale = 1024 / larger_size
+    if larger_size > max_size:
+        scale = max_size / larger_size
         image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
 
-    image_tensor = torch.tensor(image, dtype=torch.float32, device=torch.device('cuda')).permute(2, 0, 1) / 255
-    output = model.infer(image_tensor, resolution_level=9, apply_mask=True)
-    points, depth, mask = output['points'].cpu().numpy(), output['depth'].cpu().numpy(), output['mask'].cpu().numpy()
+    height, width = image.shape[:2]
+
+    output = run_with_gpu(image)
+    points, depth, mask = output['points'], output['depth'], output['mask']
 
     if remove_edge:
         mask = mask & ~utils3d.numpy.depth_edge(depth, mask=mask, rtol=0.02)
     mask = mask & (depth > 0)
 
-    _, faces, indices = utils3d.numpy.image_mesh(width=image.shape[1], height=image.shape[0], mask=mask)
-    faces = utils3d.numpy.triangulate(faces)
+    faces, vertices, vertex_colors, vertex_uvs = utils3d.numpy.image_mesh(
+        points,
+        image.astype(np.float32) / 255,
+        utils3d.numpy.image_uv(width=width, height=height),
+        mask=mask & ~utils3d.numpy.depth_edge(depth, rtol=0.02, mask=mask),
+        tri=True
+    )
+    vertices, vertex_uvs = vertices * [1, -1, -1], vertex_uvs * [1, -1] + [0, 1]
 
     tempdir = Path(tempfile.gettempdir(), 'moge')
     tempdir.mkdir(exist_ok=True)
 
     output_glb_path = Path(tempdir, f'{run_id}.glb')
     output_glb_path.parent.mkdir(exist_ok=True)
-    tempfile.TemporaryFile()
-    utils3d.io.write_glb(
-        output_glb_path,
-        vertices=points.reshape(-1, 3)[indices] * [-1, -1, 1], 
-        faces=faces,
-        vertex_colors=image.reshape(-1, 3)[indices] / 255,
-    )
+    trimesh.Trimesh(
+        vertices=vertices * [-1, 1, -1],    # No idea why Gradio 3D Viewer' default camera is flipped
+        faces=faces, 
+        visual = trimesh.visual.texture.TextureVisuals(
+            uv=vertex_uvs, 
+            material=trimesh.visual.material.PBRMaterial(
+                baseColorTexture=Image.fromarray(image),
+                metallicFactor=0.5,
+                roughnessFactor=1.0
+            )
+        ),
+        process=False
+    ).export(output_glb_path)
 
     output_ply_path = Path(tempdir, f'{run_id}.ply')
     output_ply_path.parent.mkdir(exist_ok=True)
-    utils3d.io.write_ply(
-        output_ply_path,
-        vertices=points.reshape(-1, 3)[indices] * [-1, -1, 1], 
-        faces=faces,
-        vertex_colors=image.reshape(-1, 3)[indices] / 255,
-    )
+    trimesh.Trimesh(
+        vertices=vertices, 
+        faces=faces, 
+        vertex_colors=vertex_colors,
+        process=False
+    ).export(output_ply_path)
 
     colorized_depth = colorize_depth(depth)
 
@@ -84,15 +109,16 @@ def run(image: np.ndarray, remove_edge: bool = True):
 
 
 DESCRIPTION = """
-Turns 2D images into 3D point maps with MoGe
+## Turn a 2D image into a 3D point map with [MoGe](https://wangrc.site/MoGePage/)
 
 NOTE: 
-* If the image is too large (> 1024px), it will be resized accordingly.
+* The maximum size is set to 800px for efficiency purpose. Oversized images will be downsampled.
 * The color in the 3D viewer may look dark due to rendering of 3D viewer. You may download the 3D model as .glb or .ply file to view it in other 3D viewers.
 """
 
-if __name__ == '__main__':
-    
+@click.command()
+@click.option('--share', is_flag=True, help='Whether to run the app in shared mode.')
+def main(share: bool):
     gr.Interface(
         fn=run,
         inputs=[
@@ -109,4 +135,8 @@ if __name__ == '__main__':
         clear_btn=None,
         allow_flagging="never",
         theme=gr.themes.Soft()
-    ).launch()
+    ).launch(share=share)
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -1,6 +1,5 @@
 opencv-python
-plyfile
-pygltflib
-transformers
-scikit-learn
-matplotlib
+scipy
+matplotlib
+huggingface_hub
+trimesh