import os os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' import sys from pathlib import Path if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: sys.path.insert(0, _package_root) import time import tempfile import itertools from typing import * import atexit from concurrent.futures import ThreadPoolExecutor import shutil import click TEMP_DIR = Path(tempfile.gettempdir(), 'moge') @click.command(help='Web demo') @click.option('--share', is_flag=True, help='Whether to run the app in shared mode.') @click.option('--pretrained', 'pretrained_model_name_or_path', default=None, help='Pretrained model name or path. Optional for v1/v2 and required for v3.') @click.option('--version', 'model_version', type=click.Choice(['v1', 'v2', 'v3']), default='v3', show_default=True, help='The version of the model.') @click.option('--fp16/--fp32', 'use_fp16', default=True, help='Whether to use fp16 or fp32 inference.') def main(share: bool, pretrained_model_name_or_path: Optional[str], model_version: str, use_fp16: bool): print("Import modules...") # Lazy import import cv2 import torch import numpy as np import trimesh import trimesh.visual from PIL import Image import gradio as gr from moge.utils.gradio_3d_viewer import DepthMap3DViewer from starlette.middleware import Middleware from starlette.middleware.gzip import GZipMiddleware try: import spaces # This is for deployment at huggingface.co/spaces HUGGINGFACE_SPACES_INSTALLED = True except ImportError: HUGGINGFACE_SPACES_INSTALLED = False import flex_gemm flex_gemm.config.AUTOTUNE_MODE = 'never' # Disable flex_gemm auto-tuning to avoid latency for the first inference on GPU. try: import utils3d_moge as utils3d except ImportError: import utils3d from moge.utils.io import write_normal from moge.utils.vis import colorize_depth, colorize_normal from moge.model import import_model_class_by_version from moge.utils.geometry_numpy import depth_occlusion_edge_numpy from moge.utils.tools import timeit print("Load model...") if pretrained_model_name_or_path is None: default_pretrained_models = { 'v1': 'Ruicheng/moge-vitl', 'v2': 'Ruicheng/moge-2-vitl-normal', 'v3': 'Ruicheng/moge-3-vitl' } pretrained_model_name_or_path = default_pretrained_models[model_version] model = import_model_class_by_version(model_version).from_pretrained(pretrained_model_name_or_path).cuda().eval() thread_pool_executor = ThreadPoolExecutor(max_workers=1) TEMP_DIR.mkdir(exist_ok=True) def delete_later(path: Union[str, os.PathLike], delay: int = 300): def _delete(): try: os.remove(path) except FileNotFoundError: pass def _wait_and_delete(): time.sleep(delay) _delete() thread_pool_executor.submit(_wait_and_delete) atexit.register(_delete) # Inference on GPU. # The default ZeroGPU budget of 60s is not enough for the first request, which # also pays for the Triton kernel compilation inside flex_gemm. @(spaces.GPU(duration=120) if HUGGINGFACE_SPACES_INSTALLED else lambda x: x) def run_with_gpu(image: np.ndarray, resolution_level: int, apply_mask: bool, refine_steps: int) -> Dict[str, np.ndarray]: image_tensor = torch.tensor(image, dtype=torch.float32, device=torch.device('cuda')).permute(2, 0, 1) / 255 infer_kwargs = { 'apply_mask': apply_mask, 'resolution_level': resolution_level, 'use_fp16': use_fp16, } if model_version == 'v3': infer_kwargs['refine_steps'] = refine_steps output = model.infer(image_tensor, **infer_kwargs) output = {k: v.cpu().numpy() for k, v in output.items() if isinstance(v, torch.Tensor)} return output # Full inference pipeline def run(image: np.ndarray, max_size: int = 1024, resolution_level: str = 'High', apply_mask: bool = True, remove_edge: bool = True, refine_steps: int = 3, request: gr.Request = None): larger_size = max(image.shape[:2]) if larger_size > max_size: scale = max_size / larger_size image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA) height, width = image.shape[:2] resolution_level_int = {'Low': 0, 'Medium': 5, 'High': 9, 'Ultra': 30}.get(resolution_level, 9) output = run_with_gpu(image, resolution_level_int, apply_mask, refine_steps) points, depth, mask = output['points'], output['depth'], output['mask'] normal = output.get('normal') if remove_edge: edge = utils3d.np.depth_map_edge(depth, ltol=0.01) mask_cleaned = mask & ~edge else: mask_cleaned = mask results = { **output, 'mask_cleaned': mask_cleaned, 'image': image } # depth & normal visualization depth_vis = colorize_depth(depth) normal_vis = colorize_normal(normal) if normal is not None else None mask_vis = mask_cleaned.astype(np.uint8) * 255 # mesh & pointcloud faces, vertices, vertex_colors, vertex_uvs = utils3d.np.build_mesh_from_map( points, image.astype(np.float32) / 255, utils3d.np.uv_map((height, width)), mask=mask_cleaned, tri=True ) vertices = vertices * np.array([1, -1, -1], dtype=np.float32) vertex_uvs = vertex_uvs * np.array([1, -1], dtype=np.float32) + np.array([0, 1], dtype=np.float32) TEMP_DIR.mkdir(exist_ok=True) output_path = Path(TEMP_DIR, request.session_hash) shutil.rmtree(output_path, ignore_errors=True) output_path.mkdir(exist_ok=True, parents=True) trimesh.Trimesh( vertices=vertices * np.array([-1, 1, -1], dtype=np.float32), 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_path / 'mesh.glb') trimesh.Trimesh( vertices=vertices, faces=faces, vertex_colors=vertex_colors, process=False ).export(output_path / 'mesh.ply') trimesh.PointCloud( vertices=vertices, colors=vertex_colors, ).export(output_path / 'pointcloud.glb') trimesh.PointCloud( vertices=vertices, colors=vertex_colors, ).export(output_path / 'pointcloud.ply') cv2.imwrite(str(output_path / 'depth.exr'), depth.astype(np.float32), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) cv2.imwrite(str(output_path / 'points.exr'), cv2.cvtColor(points.astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) files = ['mesh.glb', 'mesh.ply', 'pointcloud.glb', 'pointcloud.ply', 'depth.exr', 'points.exr'] if normal is not None: cv2.imwrite(str(output_path / 'normal.exr'), cv2.cvtColor(normal.astype(np.float32) * np.array([1, -1, -1], dtype=np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF]) files.append('normal.exr') for f in files: delete_later(output_path / f) # FOV intrinsics = results['intrinsics'] fov_x, fov_y = utils3d.np.intrinsics_to_fov(intrinsics) fov_x, fov_y = np.rad2deg([fov_x, fov_y]) viewer_html_value = point_cloud_viewer.build( depth, image, intrinsics, output_path, mask=mask_cleaned, ) for path in point_cloud_viewer.payload_paths(output_path): delete_later(path) # messages viewer_message = f'**Note:** Inference has been completed. The point cloud is streamed as a depth map and unprojected in your browser.' if resolution_level != 'Ultra': depth_message = f'**Note:** Want sharper depth map? Try increasing the `maximum image size` and setting the `inference resolution level` to `Ultra` in the settings.' else: depth_message = "" return ( results, depth_vis, normal_vis, mask_vis, viewer_html_value, [(output_path / f).as_posix() for f in files], f'**Horizontal FOV: {fov_x:.1f}°. Vertical FOV: {fov_y:.1f}°**', viewer_message, depth_message ) def reset_measure(results: Dict[str, np.ndarray]): return [results['image'], [], ""] def measure(results: Dict[str, np.ndarray], measure_points: List[Tuple[int, int]], event: gr.SelectData): point2d = event.index[0], event.index[1] measure_points.append(point2d) image = results['image'].copy() for p in measure_points: image = cv2.circle(image, p, radius=5, color=(255, 0, 0), thickness=2) depth_text = "" for i, p in enumerate(measure_points): d = results['depth'][p[1], p[0]] depth_text += f"**P{i + 1} depth: {d:.2f}m.** " if len(measure_points) == 2: point1, point2 = measure_points image = cv2.line(image, point1, point2, color=(255, 0, 0), thickness=2) distance = np.linalg.norm(results['points'][point1[1], point1[0]] - results['points'][point2[1], point2[0]]) measure_points = [] distance_text = f"**Distance: {distance:.2f}m**" text = depth_text + distance_text return [image, measure_points, text] else: return [image, measure_points, depth_text] print("Create Gradio app...") model_names = {'v1': 'MoGe-1', 'v2': 'MoGe-2', 'v3': 'MoGe-3'} model_urls = {'v1': 'https://wangrc.site/MoGePage/', 'v2': 'https://wangrc.site/MoGe2Page/', 'v3': 'https://qft-333.github.io/moge3page/'} model_name = model_names[model_version] with gr.Blocks() as demo: gr.Markdown( f"## Turn a 2D image into a 3D point map with [{model_name}]({model_urls[model_version]})\n Model: {pretrained_model_name_or_path}" ) results = gr.State(value=None) measure_points = gr.State(value=[]) with gr.Row(): with gr.Column(): input_image = gr.Image(type="numpy", image_mode="RGB", label="Input Image") with gr.Accordion(label="Settings", open=False): max_size_input = gr.Number(value=1024, label="Maximum Image Size", precision=0, minimum=256, maximum=4096) refine_steps = gr.Number(value=3 if hasattr(model, 'refiner') else 0, label="Refine Steps", precision=0, minimum=0, maximum=5, visible=hasattr(model, 'refiner')) resolution_level = gr.Dropdown(['Low', 'Medium', 'High', 'Ultra'], label="Inference Resolution Level", value='High') apply_mask = gr.Checkbox(value=True, label="Apply mask") remove_edges = gr.Checkbox(value=True, label="Remove edges") submit_btn = gr.Button("Submit") with gr.Column(): with gr.Tabs(): with gr.Tab("3D View"): viewer_message = gr.Markdown("") point_cloud_viewer = DepthMap3DViewer() fov = gr.Markdown() with gr.Tab("Depth"): depth_message = gr.Markdown("") depth_map = gr.Image(type="numpy", label="Colorized Depth Map", format='png', interactive=False) with gr.Tab("Normal", visible=hasattr(model, 'normal_head')): normal_map = gr.Image(type="numpy", label="Normal Map", format='png', interactive=False) with gr.Tab("Mask"): mask_map = gr.Image(type="numpy", label="Mask", format='png', interactive=False) with gr.Tab("Measure", interactive=hasattr(model, 'scale_head')): gr.Markdown("### Click on the image to measure the distance between two points. \n" "**Note:** Metric scale is most reliable for typical indoor or street scenes, and may degrade for contents unfamiliar to the model (e.g., stylized or close-up images).") measure_image = gr.Image(type="numpy", show_label=False, format='webp', interactive=False, sources=[]) gr.Markdown("Click on the image to measure the distance between two points.") measure_text = gr.Markdown("") with gr.Tab("Download"): files = gr.File(type='filepath', label="Output Files") example_image_dir = Path(_package_root, 'example_images') if example_image_dir.exists(): example_image_paths = sorted(list(itertools.chain(*[example_image_dir.glob(f'*.{ext}') for ext in ['jpg', 'png', 'jpeg', 'JPG', 'PNG', 'JPEG']]))) examples = gr.Examples( examples = example_image_paths, inputs=input_image, label="Examples" ) submit_btn.click( fn=lambda: [None, None, None, None, "", None, "", "", ""], outputs=[results, depth_map, normal_map, mask_map, point_cloud_viewer, files, fov, viewer_message, depth_message] ).then( fn=run, inputs=[input_image, max_size_input, resolution_level, apply_mask, remove_edges, refine_steps], outputs=[results, depth_map, normal_map, mask_map, point_cloud_viewer, files, fov, viewer_message, depth_message] ).then( fn=reset_measure, inputs=[results], outputs=[measure_image, measure_points, measure_text] ) measure_image.select( fn=measure, inputs=[results, measure_points], outputs=[measure_image, measure_points, measure_text] ) demo.launch( share=share, allowed_paths=[str(TEMP_DIR)], app_kwargs={ "middleware": [ Middleware( GZipMiddleware, minimum_size=1024, compresslevel=6, ) ] }, ) if __name__ == '__main__': main()