import warnings
from typing import List, Optional, Union
import torch
from pytorch3d.io import IO
from pytorch3d.io import load_objs_as_meshes as _load_objs_as_meshes
from pytorch3d.io import save_obj
from pytorch3d.renderer import TexturesUV, TexturesVertex
from pytorch3d.structures import (
Meshes,
Pointclouds,
join_meshes_as_batch,
join_meshes_as_scene,
padded_to_list,
)
from .path_utils import prepare_output_path
def join_batch_meshes_as_scene(
meshes: List[Meshes],
include_textures: bool = True,
) -> Meshes:
"""Join `meshes` as a scene each batch. Only for Pytorch3D `meshes`. The
Meshes must share the same batch size, and topology could be different.
They must all be on the same device. If `include_textures` is true, the
textures should be the same type, all be None is not accepted. If
`include_textures` is False, textures are ignored. The return meshes will
have no textures.
Args:
meshes (List[Meshes]): A `list` of `Meshes` with the same batches.
Required.
include_textures: (bool) whether to try to join the textures.
Returns:
New Meshes which has join different Meshes by each batch.
"""
for mesh in meshes:
mesh._verts_list = padded_to_list(mesh.verts_padded(),
mesh.num_verts_per_mesh().tolist())
num_scene_size = len(meshes)
num_batch_size = len(meshes[0])
for i in range(num_scene_size):
assert len(
meshes[i]
) == num_batch_size, 'Please make sure that the Meshes all have'
'the same batch size.'
meshes_all = []
for j in range(num_batch_size):
meshes_batch = []
for i in range(num_scene_size):
meshes_batch.append(meshes[i][j])
meshes_all.append(join_meshes_as_scene(meshes_batch, include_textures))
meshes_final = join_meshes_as_batch(meshes_all, include_textures)
return meshes_final
def mesh_to_pointcloud_vc(
meshes: Meshes,
include_textures: bool = True,
alpha: float = 1.0,
) -> Pointclouds:
"""Convert PyTorch3D vertex color `Meshes` to `PointClouds`.
Args:
meshes (Meshes): input meshes.
include_textures (bool, optional): Whether include colors.
Require the texture of input meshes is vertex color.
Defaults to True.
alpha (float, optional): transparency.
Defaults to 1.0.
Returns:
Pointclouds: output pointclouds.
"""
assert isinstance(
meshes.textures,
TexturesVertex), 'textures of input meshes should be `TexturesVertex`'
vertices = meshes.verts_padded()
if include_textures:
verts_rgb = meshes.textures.verts_features_padded()
verts_rgba = torch.cat(
[verts_rgb,
torch.ones_like(verts_rgb)[..., 0:1] * alpha], dim=-1)
else:
verts_rgba = None
pointclouds = Pointclouds(points=vertices, features=verts_rgba)
return pointclouds
def texture_uv2vc(meshes: Meshes) -> Meshes:
"""Convert a Pytorch3D meshes's textures from TexturesUV to TexturesVertex.
Args:
meshes (Meshes): input Meshes.
Returns:
Meshes: converted Meshes.
"""
assert isinstance(meshes.textures, TexturesUV)
device = meshes.device
vert_uv = meshes.textures.verts_uvs_padded()
batch_size = vert_uv.shape[0]
verts_features = []
num_verts = meshes.verts_padded().shape[1]
for index in range(batch_size):
face_uv = vert_uv[index][meshes.textures.faces_uvs_padded()
[index].view(-1)]
img = meshes.textures._maps_padded[index]
width, height, _ = img.shape
face_uv = face_uv * torch.Tensor([width - 1, height - 1
]).long().to(device)
face_uv[:, 0] = torch.clip(face_uv[:, 0], 0, width - 1)
face_uv[:, 1] = torch.clip(face_uv[:, 1], 0, height - 1)
face_uv = face_uv.long()
faces = meshes.faces_padded()
verts_rgb = torch.zeros(1, num_verts, 3).to(device)
verts_rgb[:, faces.view(-1)] = img[height - 1 - face_uv[:, 1],
face_uv[:, 0]]
verts_features.append(verts_rgb)
verts_features = torch.cat(verts_features)
meshes = meshes.clone()
meshes.textures = TexturesVertex(verts_features)
return meshes
def load_objs_as_meshes(files: List[str],
device: Optional[Union[torch.device, str]] = None,
load_textures: bool = True,
**kwargs) -> Meshes:
if not isinstance(files, list):
files = [files]
return _load_objs_as_meshes(files=files,
device=device,
load_textures=load_textures,
**kwargs)
def load_plys_as_meshes(
files: List[str],
device: Optional[Union[torch.device, str]] = None,
load_textures: bool = True,
) -> Meshes:
writer = IO()
meshes = []
if not isinstance(files, list):
files = [files]
for idx in range(len(files)):
assert files[idx].endswith('.ply'), 'Please input .ply files.'
mesh = writer.load_mesh(path=files[idx],
include_textures=load_textures,
device=device)
meshes.append(mesh)
meshes = join_meshes_as_batch(meshes, include_textures=load_textures)
return meshes
def save_meshes_as_plys(files: List[str],
meshes: Meshes = None,
verts: torch.Tensor = None,
faces: torch.Tensor = None,
verts_rgb: torch.Tensor = None) -> None:
"""Save meshes as .ply files. Mainly for vertex color meshes.
Args:
files (List[str]): Output .ply file list.
meshes (Meshes, optional): higher priority than
(verts & faces & verts_rgb). Defaults to None.
verts (torch.Tensor, optional): lower priority than meshes.
Defaults to None.
faces (torch.Tensor, optional): lower priority than meshes.
Defaults to None.
verts_rgb (torch.Tensor, optional): lower priority than meshes.
Defaults to None.
"""
if meshes is None:
assert verts is not None and faces is not None, 'Not mesh input.'
meshes = Meshes(
verts=verts,
faces=faces,
textures=TexturesVertex(
verts_features=verts_rgb) if verts_rgb is not None else None)
else:
if verts is not None or faces is not None or verts_rgb is not None:
warnings.warn('Redundant input, will use meshes only.')
assert files is not None
if not isinstance(files, list):
files = [files]
assert len(files) >= len(meshes), 'Not enough output files.'
writer = IO()
for idx in range(len(meshes)):
assert files[idx].endswith('.ply'), 'Please save as .ply files.'
writer.save_mesh(meshes[idx],
files[idx],
colors_as_uint8=True,
binary=False)
def save_meshes_as_objs(files: List[str], meshes: Meshes = None) -> None:
"""Save meshes as .obj files. Pytorch3D will not save vertex color for.
.obj, please use `save_meshes_as_plys`.
Args:
files (List[str]): Output .obj file list.
meshes (Meshes, optional):
Defaults to None.
"""
if not isinstance(files, list):
files = [files]
assert len(files) >= len(meshes), 'Not enough output files.'
for idx in range(len(meshes)):
prepare_output_path(files[idx],
allowed_suffix=['.obj'],
path_type='file'), 'Please save as .obj files.'
if isinstance(meshes.textures, TexturesUV):
verts_uvs = meshes.textures.verts_uvs_padded()[idx]
faces_uvs = meshes.textures.faces_uvs_padded()[idx]
texture_map = meshes.textures.maps_padded()[idx]
else:
verts_uvs = None
faces_uvs = None
texture_map = None
save_obj(f=files[idx],
verts=meshes.verts_padded()[idx],
faces=meshes.faces_padded()[idx],
verts_uvs=verts_uvs,
faces_uvs=faces_uvs,
texture_map=texture_map)