File size: 4,429 Bytes
ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 239102c ec0c8fa 728b726 ec0c8fa 728b726 239102c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
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)
mask = mask & (depth > 0)
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)
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)
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() |