171 lines
8.8 KiB
Python
171 lines
8.8 KiB
Python
import os
|
|
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
|
from pathlib import Path
|
|
import sys
|
|
if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path:
|
|
sys.path.insert(0, _package_root)
|
|
from typing import *
|
|
import itertools
|
|
import json
|
|
import warnings
|
|
|
|
import click
|
|
|
|
|
|
@click.command(help='Inference script')
|
|
@click.option('--input', '-i', 'input_path', type=click.Path(exists=True), help='Input image or folder path. "jpg" and "png" are supported.')
|
|
@click.option('--fov_x', 'fov_x_', type=float, default=None, help='If camera parameters are known, set the horizontal field of view in degrees. Otherwise, MoGe will estimate it.')
|
|
@click.option('--output', '-o', 'output_path', default='./output', type=click.Path(), help='Output folder path')
|
|
@click.option('--pretrained', 'pretrained_model_name_or_path', type=str, default=None, help='Pretrained model name or path. If not provided, the corresponding default model will be chosen.')
|
|
@click.option('--version', 'model_version', type=click.Choice(['v1', 'v2']), default='v2', help='Model version. Defaults to "v2"')
|
|
@click.option('--device', 'device_name', type=str, default='cuda', help='Device name (e.g. "cuda", "cuda:0", "cpu"). Defaults to "cuda"')
|
|
@click.option('--fp16', 'use_fp16', is_flag=True, help='Use fp16 precision for much faster inference.')
|
|
@click.option('--resize', 'resize_to', type=int, default=None, help='Resize the image(s) & output maps to a specific size. Defaults to None (no resizing).')
|
|
@click.option('--resolution_level', type=int, default=9, help='An integer [0-9] for the resolution level for inference. \
|
|
Higher value means more tokens and the finer details will be captured, but inference can be slower. \
|
|
Defaults to 9. Note that it is irrelevant to the output size, which is always the same as the input size. \
|
|
`resolution_level` actually controls `num_tokens`. See `num_tokens` for more details.')
|
|
@click.option('--num_tokens', type=int, default=None, help='number of tokens used for inference. A integer in the (suggested) range of `[1200, 2500]`. \
|
|
`resolution_level` will be ignored if `num_tokens` is provided. Default: None')
|
|
@click.option('--threshold', type=float, default=0.04, help='Threshold for removing edges. Defaults to 0.01. Smaller value removes more edges. "inf" means no thresholding.')
|
|
@click.option('--maps', 'save_maps_', is_flag=True, help='Whether to save the output maps (image, point map, depth map, normal map, mask) and fov.')
|
|
@click.option('--glb', 'save_glb_', is_flag=True, help='Whether to save the output as a.glb file. The color will be saved as a texture.')
|
|
@click.option('--ply', 'save_ply_', is_flag=True, help='Whether to save the output as a.ply file. The color will be saved as vertex colors.')
|
|
@click.option('--show', 'show', is_flag=True, help='Whether show the output in a window. Note that this requires pyglet<2 installed as required by trimesh.')
|
|
def main(
|
|
input_path: str,
|
|
fov_x_: float,
|
|
output_path: str,
|
|
pretrained_model_name_or_path: str,
|
|
model_version: str,
|
|
device_name: str,
|
|
use_fp16: bool,
|
|
resize_to: int,
|
|
resolution_level: int,
|
|
num_tokens: int,
|
|
threshold: float,
|
|
save_maps_: bool,
|
|
save_glb_: bool,
|
|
save_ply_: bool,
|
|
show: bool,
|
|
):
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
import click
|
|
|
|
from moge.model import import_model_class_by_version
|
|
from moge.utils.io import save_glb, save_ply
|
|
from moge.utils.vis import colorize_depth, colorize_normal
|
|
from moge.utils.geometry_numpy import depth_occlusion_edge_numpy
|
|
import utils3d
|
|
|
|
device = torch.device(device_name)
|
|
|
|
include_suffices = ['jpg', 'png', 'jpeg', 'JPG', 'PNG', 'JPEG']
|
|
if Path(input_path).is_dir():
|
|
image_paths = sorted(itertools.chain(*(Path(input_path).rglob(f'*.{suffix}') for suffix in include_suffices)))
|
|
else:
|
|
image_paths = [Path(input_path)]
|
|
|
|
if len(image_paths) == 0:
|
|
raise FileNotFoundError(f'No image files found in {input_path}')
|
|
|
|
if pretrained_model_name_or_path is None:
|
|
DEFAULT_PRETRAINED_MODEL_FOR_EACH_VERSION = {
|
|
"v1": "Ruicheng/moge-vitl",
|
|
"v2": "Ruicheng/moge-2-vitl-normal",
|
|
}
|
|
pretrained_model_name_or_path = DEFAULT_PRETRAINED_MODEL_FOR_EACH_VERSION[model_version]
|
|
model = import_model_class_by_version(model_version).from_pretrained(pretrained_model_name_or_path).to(device).eval()
|
|
if use_fp16:
|
|
model.half()
|
|
|
|
if not any([save_maps_, save_glb_, save_ply_]):
|
|
warnings.warn('No output format specified. Defaults to saving all. Please use "--maps", "--glb", or "--ply" to specify the output.')
|
|
save_maps_ = save_glb_ = save_ply_ = True
|
|
|
|
for image_path in (pbar := tqdm(image_paths, desc='Inference', disable=len(image_paths) <= 1)):
|
|
if not image_path.exists():
|
|
raise FileNotFoundError(f'File {image_path} does not exist.')
|
|
image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
|
|
height, width = image.shape[:2]
|
|
if resize_to is not None:
|
|
height, width = min(resize_to, int(resize_to * height / width)), min(resize_to, int(resize_to * width / height))
|
|
image = cv2.resize(image, (width, height), cv2.INTER_AREA)
|
|
image_tensor = torch.tensor(image / 255, dtype=torch.float32, device=device).permute(2, 0, 1)
|
|
|
|
# Inference
|
|
output = model.infer(image_tensor, fov_x=fov_x_, resolution_level=resolution_level, num_tokens=num_tokens, use_fp16=use_fp16)
|
|
points, depth, mask, intrinsics = output['points'].cpu().numpy(), output['depth'].cpu().numpy(), output['mask'].cpu().numpy(), output['intrinsics'].cpu().numpy()
|
|
normal = output['normal'].cpu().numpy() if 'normal' in output else None
|
|
|
|
save_path = Path(output_path, image_path.relative_to(input_path).parent, image_path.stem)
|
|
save_path.mkdir(exist_ok=True, parents=True)
|
|
|
|
# Save images / maps
|
|
if save_maps_:
|
|
cv2.imwrite(str(save_path / 'image.jpg'), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
|
|
cv2.imwrite(str(save_path / 'depth_vis.png'), cv2.cvtColor(colorize_depth(depth), cv2.COLOR_RGB2BGR))
|
|
cv2.imwrite(str(save_path / 'depth.exr'), depth, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT])
|
|
cv2.imwrite(str(save_path / 'mask.png'), (mask * 255).astype(np.uint8))
|
|
cv2.imwrite(str(save_path / 'points.exr'), cv2.cvtColor(points, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT])
|
|
if normal is not None:
|
|
cv2.imwrite(str(save_path / 'normal.png'), cv2.cvtColor(colorize_normal(normal), cv2.COLOR_RGB2BGR))
|
|
fov_x, fov_y = utils3d.np.intrinsics_to_fov(intrinsics)
|
|
with open(save_path / 'fov.json', 'w') as f:
|
|
json.dump({
|
|
'fov_x': round(float(np.rad2deg(fov_x)), 2),
|
|
'fov_y': round(float(np.rad2deg(fov_y)), 2),
|
|
}, f)
|
|
|
|
# Export mesh & visulization
|
|
if save_glb_ or save_ply_ or show:
|
|
mask_cleaned = mask & ~utils3d.np.depth_map_edge(depth, rtol=threshold)
|
|
if normal is None:
|
|
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
|
|
)
|
|
vertex_normals = None
|
|
else:
|
|
faces, vertices, vertex_colors, vertex_uvs, vertex_normals = utils3d.np.build_mesh_from_map(
|
|
points,
|
|
image.astype(np.float32) / 255,
|
|
utils3d.np.uv_map(height, width),
|
|
normal,
|
|
mask=mask_cleaned,
|
|
tri=True
|
|
)
|
|
# When exporting the model, follow the OpenGL coordinate conventions:
|
|
# - world coordinate system: x right, y up, z backward.
|
|
# - texture coordinate system: (0, 0) for left-bottom, (1, 1) for right-top.
|
|
vertices, vertex_uvs = vertices * [1, -1, -1], vertex_uvs * [1, -1] + [0, 1]
|
|
if normal is not None:
|
|
vertex_normals = vertex_normals * [1, -1, -1]
|
|
|
|
if save_glb_:
|
|
save_glb(save_path / 'mesh.glb', vertices, faces, vertex_uvs, image, vertex_normals)
|
|
|
|
if save_ply_:
|
|
save_ply(save_path / 'pointcloud.ply', vertices, np.zeros((0, 3), dtype=np.int32), vertex_colors, vertex_normals)
|
|
|
|
if show:
|
|
import trimesh
|
|
trimesh.Trimesh(
|
|
vertices=vertices,
|
|
vertex_colors=vertex_colors,
|
|
vertex_normals=vertex_normals,
|
|
faces=faces,
|
|
process=False
|
|
).show()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|