|
import os |
|
import time |
|
from pathlib import Path |
|
import uuid |
|
import tempfile |
|
from typing import Union |
|
import spaces |
|
import atexit |
|
from concurrent.futures import ThreadPoolExecutor |
|
|
|
import gradio as gr |
|
import cv2 |
|
import torch |
|
import numpy as np |
|
|
|
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(image: np.ndarray, remove_edge: bool = True): |
|
run_id = str(uuid.uuid4()) |
|
|
|
larger_size = max(image.shape[:2]) |
|
if larger_size > 1024: |
|
scale = 1024 / 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() |
|
|
|
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) |
|
|
|
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, |
|
) |
|
|
|
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, |
|
) |
|
|
|
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 = """ |
|
Turns 2D images into 3D point maps with MoGe |
|
|
|
NOTE: |
|
* If the image is too large (> 1024px), it will be resized accordingly. |
|
* 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__': |
|
|
|
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() |
