Kiss3DGen / models /ISOMER /scripts /proj_commands.py
JiantaoLin
initial commit
98bebfc
import numpy as np
import torch
from PIL import Image
from pytorch3d.renderer import (
TexturesVertex,
)
from .project_mesh import (
get_cameras_list_azi_ele,
multiview_color_projection
)
from .utils import save_py3dmesh_with_trimesh_fast
def projection(meshes,
img_list,
weights,
azimuths,
elevations,
projection_type='orthographic',
auto_center=True,
resolution=1024,
fovy=None,
radius=None,
ortho_dist=1.1,
scale_factor=1.0,
save_glb_addr=None,
scale_verts=True,
complete_unseen=True,
below_confidence_strategy="smooth"
):
assert len(img_list) == len(azimuths) == len(elevations) == len(weights), f"len(img_list) ({len(img_list)}) != len(azimuths) ({len(azimuths)}) != len(elevations) ({len(elevations)}) != len(weights) ({len(weights)})"
projection_types = ['perspective', 'orthographic']
assert projection_type in projection_types, f"projection_type ({projection_type}) should be one of {projection_types}"
if auto_center:
verts = meshes.verts_packed()
max_bb = (verts - 0).max(0)[0]
min_bb = (verts - 0).min(0)[0]
scale = (max_bb - min_bb).max() / 2
center = (max_bb + min_bb) / 2
meshes.offset_verts_(-center)
if scale_verts:
meshes.scale_verts_((scale_factor / float(scale)))
elif scale_verts:
meshes.scale_verts_((scale_factor))
if projection_type == 'perspective':
assert fovy is not None and radius is not None, f"fovy ({fovy}) and radius ({radius}) should not be None when projection_type is 'perspective'"
cameras = get_cameras_list_azi_ele(azimuths, elevations, fov_in_degrees=fovy,device="cuda", dist=radius, cam_type='fov')
elif projection_type == 'orthographic':
cameras = get_cameras_list_azi_ele(azimuths, elevations, fov_in_degrees=fovy, device="cuda", focal=2/1.35, dist=ortho_dist, cam_type='orthographic')
num_meshes = len(meshes)
num_verts_per_mesh = meshes.verts_packed().shape[0] // num_meshes
black_texture = torch.zeros((num_meshes, num_verts_per_mesh, 3), device="cuda")
textures = TexturesVertex(verts_features=black_texture)
meshes.textures = textures
proj_mesh = multiview_color_projection(meshes, img_list, cameras, weights=weights, eps=0.05, resolution=resolution, device="cuda", reweight_with_cosangle="square", use_alpha=True, confidence_threshold=0.1, complete_unseen=complete_unseen, below_confidence_strategy=below_confidence_strategy)
if save_glb_addr is not None:
save_py3dmesh_with_trimesh_fast(proj_mesh, save_glb_addr)