|
import os |
|
import time |
|
from pathlib import Path |
|
import uuid |
|
import tempfile |
|
from typing import Union |
|
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 |
|
import utils3d |
|
|
|
model = MoGeModel.from_pretrained('Ruicheng/moge-vitl').cuda().eval() |
|
thread_pool_executor = ThreadPoolExecutor(max_workers=1) |
|
|
|
|
|
def delete_later(path: Union[str, os.PathLike], delay: int = 300): |
|
def _delete(): |
|
try: |
|
os.remove(path) |
|
except: |
|
pass |
|
def _wait_and_delete(): |
|
time.sleep(delay) |
|
_delete(path) |
|
thread_pool_executor.submit(_wait_and_delete) |
|
atexit.register(_delete) |
|
|
|
|
|
@spaces.GPU |
|
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 > 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] |
|
|
|
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) |
|
|
|
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, |
|
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) |
|
trimesh.Trimesh( |
|
vertices=vertices * [-1, 1, -1], |
|
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) |
|
trimesh.Trimesh( |
|
vertices=vertices, |
|
faces=faces, |
|
vertex_colors=vertex_colors, |
|
process=False |
|
).export(output_ply_path) |
|
|
|
colorized_depth = colorize_depth(depth) |
|
|
|
delete_later(output_glb_path, delay=300) |
|
delete_later(output_ply_path, delay=300) |
|
|
|
return colorized_depth, output_glb_path, output_ply_path.as_posix() |
|
|
|
|
|
DESCRIPTION = """ |
|
## Turn a 2D image into a 3D point map with [MoGe](https://wangrc.site/MoGePage/) |
|
|
|
NOTE: |
|
* 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. |
|
""" |
|
|
|
@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=[ |
|
gr.Image(type="numpy", image_mode="RGB"), |
|
gr.Checkbox(True, label="Remove edges"), |
|
], |
|
outputs=[ |
|
gr.Image(type="numpy", label="Depth map (colorized)"), |
|
gr.Model3D(display_mode="solid", clear_color=[1.0, 1.0, 1.0, 1.0], label="3D Viewer"), |
|
gr.File(type="filepath", label="Download the model as .ply file"), |
|
], |
|
title=None, |
|
description=DESCRIPTION, |
|
clear_btn=None, |
|
allow_flagging="never", |
|
theme=gr.themes.Soft() |
|
).launch(share=share) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
