integrated fastmesh to hun2.1
This commit is contained in:
DenisKochetov
@@ -6,6 +6,6 @@ except ImportError:
|
||||
try:
|
||||
from .mesh_inpaint_processor_fallback import meshVerticeInpaint
|
||||
print("[WARNING] Using Python fallback mesh_inpaint_processor (slower)")
|
||||
print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat")
|
||||
print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat or compile_mesh_painter.sh")
|
||||
except ImportError:
|
||||
print("[ERROR] Neither C++ nor Python mesh_inpaint_processor found!")
|
||||
|
||||
+1
-1
@@ -6,6 +6,6 @@ except ImportError:
|
||||
try:
|
||||
from .mesh_inpaint_processor_fallback import meshVerticeInpaint
|
||||
print("[WARNING] Using Python fallback mesh_inpaint_processor (slower)")
|
||||
print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat")
|
||||
print("[INFO] To use faster C++ version, run: compile_mesh_painter.bat or compile_mesh_painter.sh")
|
||||
except ImportError:
|
||||
print("[ERROR] Neither C++ nor Python mesh_inpaint_processor found!")
|
||||
|
||||
@@ -35,8 +35,6 @@ from .validate import validate_before_save
|
||||
from .mesh import (
|
||||
Mesh,
|
||||
FastMesh,
|
||||
TrimeshLike,
|
||||
TrimeshLikeLoader,
|
||||
PointCloud,
|
||||
)
|
||||
|
||||
@@ -67,8 +65,6 @@ __all__ = [
|
||||
"validate_before_save",
|
||||
"Mesh",
|
||||
"FastMesh",
|
||||
"TrimeshLike",
|
||||
"TrimeshLikeLoader",
|
||||
"PointCloud",
|
||||
]
|
||||
|
||||
|
||||
@@ -0,0 +1,192 @@
|
||||
"""
|
||||
Utilities for exporting from pipelines to our mesh objects
|
||||
Direct functions WITHOUT trimesh dependency
|
||||
"""
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from .mesh import Mesh, FastMesh
|
||||
|
||||
|
||||
def export_to_fastmesh(mesh_output, device=None):
|
||||
"""
|
||||
Export from pipeline to FastMesh
|
||||
|
||||
Args:
|
||||
mesh_output: object with mesh_v and mesh_f from pipeline
|
||||
device: torch device
|
||||
|
||||
Returns:
|
||||
FastMesh object or list of FastMesh objects
|
||||
"""
|
||||
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if isinstance(mesh_output, list):
|
||||
outputs = []
|
||||
for mesh in mesh_output:
|
||||
if mesh is None:
|
||||
outputs.append(None)
|
||||
else:
|
||||
fastmesh = _convert_single_mesh_to_fastmesh(mesh, device)
|
||||
outputs.append(fastmesh)
|
||||
return outputs
|
||||
else:
|
||||
return _convert_single_mesh_to_fastmesh(mesh_output, device)
|
||||
|
||||
|
||||
def export_to_mesh(mesh_output, device=None):
|
||||
"""
|
||||
Export from pipeline to standard Mesh
|
||||
|
||||
Args:
|
||||
mesh_output: object with mesh_v and mesh_f from pipeline
|
||||
device: torch device
|
||||
|
||||
Returns:
|
||||
Mesh object or list of Mesh objects
|
||||
"""
|
||||
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if isinstance(mesh_output, list):
|
||||
outputs = []
|
||||
for mesh in mesh_output:
|
||||
if mesh is None:
|
||||
outputs.append(None)
|
||||
else:
|
||||
standard_mesh = _convert_single_mesh_to_mesh(mesh, device)
|
||||
outputs.append(standard_mesh)
|
||||
return outputs
|
||||
else:
|
||||
return _convert_single_mesh_to_mesh(mesh_output, device)
|
||||
|
||||
|
||||
def _convert_single_mesh_to_fastmesh(mesh_obj, device):
|
||||
"""Convert one mesh object to FastMesh"""
|
||||
|
||||
try:
|
||||
# Extract vertices and faces
|
||||
vertices = mesh_obj.mesh_v
|
||||
faces = mesh_obj.mesh_f
|
||||
|
||||
# Flip faces (as in original function)
|
||||
faces = faces[:, ::-1]
|
||||
|
||||
# Convert to numpy if needed
|
||||
if hasattr(vertices, 'cpu'):
|
||||
vertices = vertices.cpu().numpy()
|
||||
if hasattr(faces, 'cpu'):
|
||||
faces = faces.cpu().numpy()
|
||||
|
||||
vertices = np.array(vertices, dtype=np.float32)
|
||||
faces = np.array(faces, dtype=np.int32)
|
||||
|
||||
print(f"[export_to_fastmesh] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
|
||||
|
||||
fastmesh = FastMesh(device=device)
|
||||
fastmesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
||||
fastmesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
||||
|
||||
fastmesh.auto_normal()
|
||||
|
||||
fastmesh._create_empty_albedo()
|
||||
|
||||
print(f"[export_to_fastmesh] FastMesh created successfully")
|
||||
return fastmesh
|
||||
|
||||
except Exception as e:
|
||||
print(f"[export_to_fastmesh] Error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _convert_single_mesh_to_mesh(mesh_obj, device):
|
||||
"""Convert one mesh object to standard Mesh"""
|
||||
|
||||
try:
|
||||
vertices = mesh_obj.mesh_v
|
||||
faces = mesh_obj.mesh_f
|
||||
|
||||
faces = faces[:, ::-1]
|
||||
|
||||
if hasattr(vertices, 'cpu'):
|
||||
vertices = vertices.cpu().numpy()
|
||||
if hasattr(faces, 'cpu'):
|
||||
faces = faces.cpu().numpy()
|
||||
|
||||
vertices = np.array(vertices, dtype=np.float32)
|
||||
faces = np.array(faces, dtype=np.int32)
|
||||
|
||||
print(f"[export_to_mesh] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
|
||||
|
||||
mesh = Mesh(device=device)
|
||||
mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
||||
mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
||||
|
||||
mesh.auto_normal()
|
||||
|
||||
mesh._create_empty_albedo_fast()
|
||||
|
||||
print(f"[export_to_mesh] Mesh created successfully")
|
||||
return mesh
|
||||
|
||||
except Exception as e:
|
||||
print(f"[export_to_mesh] Error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def export_to_mesh_ultra_fast(mesh_output, device=None):
|
||||
"""
|
||||
Ultra-fast export to standard Mesh WITHOUT any slow operations
|
||||
"""
|
||||
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
try:
|
||||
# Process list or single object
|
||||
if isinstance(mesh_output, list):
|
||||
return [_convert_single_mesh_to_mesh_ultra_fast(mesh, device) for mesh in mesh_output]
|
||||
else:
|
||||
return _convert_single_mesh_to_mesh_ultra_fast(mesh_output, device)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[export_to_mesh_ultra_fast] Error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _convert_single_mesh_to_mesh_ultra_fast(mesh_obj, device):
|
||||
"""Convert one mesh object to standard Mesh Ultra-fast"""
|
||||
|
||||
try:
|
||||
# Extract vertices and faces
|
||||
vertices = mesh_obj.mesh_v
|
||||
faces = mesh_obj.mesh_f
|
||||
|
||||
# Flip faces (as in original function)
|
||||
faces = faces[:, ::-1]
|
||||
|
||||
# Convert to numpy if needed
|
||||
if hasattr(vertices, 'cpu'):
|
||||
vertices = vertices.cpu().numpy()
|
||||
if hasattr(faces, 'cpu'):
|
||||
faces = faces.cpu().numpy()
|
||||
|
||||
vertices = np.array(vertices, dtype=np.float32)
|
||||
faces = np.array(faces, dtype=np.int32)
|
||||
|
||||
print(f"[export_to_mesh_ultra_fast] Convert: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
|
||||
|
||||
# Create standard Mesh Minimal
|
||||
mesh = Mesh(device=device)
|
||||
mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
|
||||
mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
|
||||
|
||||
print(f"[export_to_mesh_ultra_fast] Mesh created instantly")
|
||||
return mesh
|
||||
|
||||
except Exception as e:
|
||||
print(f"[export_to_mesh_ultra_fast] Error: {e}")
|
||||
return None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,741 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from kornia.geometry.conversions import (
|
||||
quaternion_to_axis_angle,
|
||||
axis_angle_to_quaternion,
|
||||
)
|
||||
import pymeshlab as pml
|
||||
from plyfile import PlyData, PlyElement
|
||||
|
||||
from typing import Optional
|
||||
|
||||
pytorch3d_capable = True
|
||||
try:
|
||||
import pytorch3d
|
||||
from pytorch3d.ops import knn_points
|
||||
except ImportError:
|
||||
pytorch3d_capable = False
|
||||
|
||||
from .mesh import PointCloud
|
||||
from shared_utils.sh_utils import SH2RGB, RGB2SH
|
||||
|
||||
def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2):
|
||||
"""Base function to compute the face areas."""
|
||||
x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1)
|
||||
y1, y2, y3 = torch.split(face_vertices_1 - face_vertices_2, 1, dim=-1)
|
||||
|
||||
a = (x2 * y3 - x3 * y2) ** 2
|
||||
b = (x3 * y1 - x1 * y3) ** 2
|
||||
c = (x1 * y2 - x2 * y1) ** 2
|
||||
areas = torch.sqrt(a + b + c) * 0.5
|
||||
|
||||
return areas
|
||||
|
||||
|
||||
def _base_sample_points_selected_faces(face_vertices, face_features=None):
|
||||
"""Base function to sample points over selected faces, sample one point per face.
|
||||
The coordinates of the face vertices are interpolated to generate new samples.
|
||||
|
||||
Args:
|
||||
face_vertices (tuple of torch.Tensor):
|
||||
Coordinates of vertices, corresponding to selected faces to sample from.
|
||||
A tuple of 3 entries corresponding to each of the face vertices.
|
||||
Each entry is a torch.Tensor of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`.
|
||||
face_features (tuple of torch.Tensor, Optional):
|
||||
Features of face vertices, corresponding to selected faces to sample from.
|
||||
A tuple of 3 entries corresponding to each of the face vertices.
|
||||
Each entry is a torch.Tensor of shape
|
||||
:math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor, torch.Tensor):
|
||||
Sampled point coordinates of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`.
|
||||
Sampled points interpolated features of shape
|
||||
:math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`.
|
||||
If `face_vertices_features` arg is not specified, the returned interpolated features are None.
|
||||
"""
|
||||
|
||||
face_vertices0, face_vertices1, face_vertices2 = face_vertices
|
||||
|
||||
sampling_shape = tuple(int(d) for d in face_vertices0.shape[:-1]) + (1,)
|
||||
# u is proximity to middle point between v1 and v2 against v0.
|
||||
# v is proximity to v2 against v1.
|
||||
#
|
||||
# The probability density for u should be f_U(u) = 2u.
|
||||
# However, torch.rand use a uniform (f_X(x) = x) distribution,
|
||||
# so using torch.sqrt we make a change of variable to have the desired density
|
||||
# f_Y(y) = f_X(y ^ 2) * |d(y ^ 2) / dy| = 2y
|
||||
u = torch.sqrt(torch.rand(sampling_shape,
|
||||
device=face_vertices0.device,
|
||||
dtype=face_vertices0.dtype))
|
||||
|
||||
v = torch.rand(sampling_shape,
|
||||
device=face_vertices0.device,
|
||||
dtype=face_vertices0.dtype)
|
||||
w0 = 1 - u
|
||||
w1 = u * (1 - v)
|
||||
w2 = u * v
|
||||
|
||||
points = w0 * face_vertices0 + w1 * face_vertices1 + w2 * face_vertices2
|
||||
|
||||
features = None
|
||||
if face_features is not None:
|
||||
face_features0, face_features1, face_features2 = face_features
|
||||
features = w0 * face_features0 + w1 * face_features1 + \
|
||||
w2 * face_features2
|
||||
|
||||
return points, features
|
||||
|
||||
# Modified from https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L159
|
||||
def sample_points(vertices, faces, num_samples, areas=None, face_features=None):
|
||||
r"""Uniformly sample points over the surface of triangle meshes.
|
||||
|
||||
First, face on which the point is sampled is randomly selected,
|
||||
with the probability of selection being proportional to the area of the face.
|
||||
then the coordinate on the face is uniformly sampled.
|
||||
|
||||
If ``face_features`` is defined for the mesh faces,
|
||||
the sampled points will be returned with interpolated features as well,
|
||||
otherwise, no feature interpolation will occur.
|
||||
|
||||
Args:
|
||||
vertices (torch.Tensor):
|
||||
The vertices of the meshes, of shape
|
||||
:math:`(\text{batch_size}, \text{num_vertices}, 3)`.
|
||||
faces (torch.LongTensor):
|
||||
The faces of the mesh, of shape :math:`(\text{num_faces}, 3)`.
|
||||
num_samples (int):
|
||||
The number of point sampled per mesh.
|
||||
Also the number of faces sampled per mesh, and then sample a single point per face.
|
||||
areas (torch.Tensor, optional):
|
||||
The areas of each face, of shape :math:`(\text{batch_size}, \text{num_faces})`,
|
||||
can be preprocessed, for fast on-the-fly sampling,
|
||||
will be computed if None (default).
|
||||
face_features (torch.Tensor, optional):
|
||||
Per-vertex-per-face features, matching ``faces`` order,
|
||||
of shape :math:`(\text{batch_size}, \text{num_faces}, 3, \text{feature_dim})`.
|
||||
For example:
|
||||
|
||||
1. Texture uv coordinates would be of shape
|
||||
:math:`(\text{batch_size}, \text{num_faces}, 3, 2)`.
|
||||
2. RGB color values would be of shape
|
||||
:math:`(\text{batch_size}, \text{num_faces}, 3, 3)`.
|
||||
|
||||
When specified, it is used to interpolate the features for new sampled points.
|
||||
|
||||
See also:
|
||||
:func:`~kaolin.ops.mesh.index_vertices_by_faces` for conversion of features defined per vertex
|
||||
and need to be converted to per-vertex-per-face shape of :math:`(\text{num_faces}, 3)`.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor, torch.LongTensor, (optional) torch.Tensor):
|
||||
the pointclouds of shape :math:`(\text{batch_size}, \text{num_samples}, 3)`,
|
||||
and the indexes of the faces selected,
|
||||
of shape :math:`(\text{batch_size}, \text{num_samples})`.
|
||||
|
||||
If ``face_features`` arg is specified, then the interpolated features of sampled points of shape
|
||||
:math:`(\text{batch_size}, \text{num_samples}, \text{feature_dim})` are also returned.
|
||||
"""
|
||||
if faces.shape[-1] != 3:
|
||||
raise NotImplementedError("sample_points is only implemented for triangle meshes")
|
||||
faces_0, faces_1, faces_2 = torch.split(faces, 1, dim=1) # (num_faces, 3) -> tuple of (num_faces,)
|
||||
face_v_0 = torch.index_select(vertices, 1, faces_0.reshape(-1)) # (batch_size, num_faces, 3)
|
||||
face_v_1 = torch.index_select(vertices, 1, faces_1.reshape(-1)) # (batch_size, num_faces, 3)
|
||||
face_v_2 = torch.index_select(vertices, 1, faces_2.reshape(-1)) # (batch_size, num_faces, 3)
|
||||
|
||||
if areas is None:
|
||||
areas = _base_face_areas(face_v_0, face_v_1, face_v_2).squeeze(-1)
|
||||
face_dist = torch.distributions.Categorical(areas)
|
||||
face_choices = face_dist.sample([num_samples]).transpose(0, 1)
|
||||
_face_choices = face_choices.unsqueeze(-1).repeat(1, 1, 3)
|
||||
v0 = torch.gather(face_v_0, 1, _face_choices) # (batch_size, num_samples, 3)
|
||||
v1 = torch.gather(face_v_1, 1, _face_choices) # (batch_size, num_samples, 3)
|
||||
v2 = torch.gather(face_v_2, 1, _face_choices) # (batch_size, num_samples, 3)
|
||||
face_vertices_choices = (v0, v1, v2)
|
||||
|
||||
# UV coordinates are available, make sure to calculate them for sampled points as well
|
||||
face_features_choices = None
|
||||
if face_features is not None:
|
||||
feat_dim = face_features.shape[-1]
|
||||
# (num_faces, 3) -> tuple of (num_faces,)
|
||||
_face_choices = face_choices[..., None, None].repeat(1, 1, 3, feat_dim)
|
||||
face_features_choices = torch.gather(face_features, 1, _face_choices)
|
||||
face_features_choices = tuple(
|
||||
tmp_feat.squeeze(2) for tmp_feat in torch.split(face_features_choices, 1, dim=2))
|
||||
|
||||
points, point_features = _base_sample_points_selected_faces(
|
||||
face_vertices_choices, face_features_choices)
|
||||
|
||||
if point_features is not None:
|
||||
return points, face_choices, point_features
|
||||
else:
|
||||
return points, face_choices
|
||||
|
||||
def poisson_mesh_reconstruction(points, normals=None):
|
||||
# points/normals: [N, 3] np.ndarray
|
||||
|
||||
import open3d as o3d
|
||||
|
||||
pcd = o3d.geometry.PointCloud()
|
||||
pcd.points = o3d.utility.Vector3dVector(points)
|
||||
|
||||
# outlier removal
|
||||
pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=10)
|
||||
|
||||
# normals
|
||||
if normals is None:
|
||||
pcd.estimate_normals()
|
||||
else:
|
||||
pcd.normals = o3d.utility.Vector3dVector(normals[ind])
|
||||
|
||||
# visualize
|
||||
o3d.visualization.draw_geometries([pcd], point_show_normal=False)
|
||||
|
||||
mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
|
||||
pcd, depth=9
|
||||
)
|
||||
vertices_to_remove = densities < np.quantile(densities, 0.1)
|
||||
mesh.remove_vertices_by_mask(vertices_to_remove)
|
||||
|
||||
# visualize
|
||||
o3d.visualization.draw_geometries([mesh])
|
||||
|
||||
vertices = np.asarray(mesh.vertices)
|
||||
triangles = np.asarray(mesh.triangles)
|
||||
|
||||
print(
|
||||
f"[INFO] poisson mesh reconstruction: {points.shape} --> {vertices.shape} / {triangles.shape}"
|
||||
)
|
||||
|
||||
return vertices, triangles
|
||||
|
||||
|
||||
def decimate_mesh(
|
||||
verts, faces, target=5e4, backend="pymeshlab", remesh=False, optimalplacement=True, verbose=True
|
||||
):
|
||||
""" perform mesh decimation.
|
||||
|
||||
Args:
|
||||
verts (np.ndarray): mesh vertices, float [N, 3]
|
||||
faces (np.ndarray): mesh faces, int [M, 3]
|
||||
target (int): targeted number of faces
|
||||
backend (str, optional): algorithm backend, can be "pymeshlab" or "pyfqmr". Defaults to "pymeshlab".
|
||||
remesh (bool, optional): whether to remesh after decimation. Defaults to False.
|
||||
optimalplacement (bool, optional): For flat mesh, use False to prevent spikes. Defaults to True.
|
||||
verbose (bool, optional): whether to print the decimation process. Defaults to True.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray]: vertices and faces after decimation.
|
||||
"""
|
||||
|
||||
_ori_vert_shape = verts.shape
|
||||
_ori_face_shape = faces.shape
|
||||
|
||||
if backend == "pyfqmr":
|
||||
import pyfqmr
|
||||
|
||||
solver = pyfqmr.Simplify()
|
||||
solver.setMesh(verts, faces)
|
||||
solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False)
|
||||
verts, faces, normals = solver.getMesh()
|
||||
else:
|
||||
m = pml.Mesh(verts, faces)
|
||||
ms = pml.MeshSet()
|
||||
ms.add_mesh(m, "mesh") # will copy!
|
||||
|
||||
ms.meshing_merge_close_vertices()
|
||||
|
||||
# filters
|
||||
# ms.meshing_decimation_clustering(threshold=pml.PercentageValue(1))
|
||||
ms.meshing_decimation_quadric_edge_collapse(
|
||||
targetfacenum=int(target), optimalplacement=optimalplacement
|
||||
)
|
||||
|
||||
if remesh:
|
||||
# ms.apply_coord_taubin_smoothing()
|
||||
ms.meshing_isotropic_explicit_remeshing(
|
||||
iterations=3, targetlen=pml.PercentageValue(1)
|
||||
)
|
||||
|
||||
# extract mesh
|
||||
m = ms.current_mesh()
|
||||
m.compact()
|
||||
verts = m.vertex_matrix()
|
||||
faces = m.face_matrix()
|
||||
|
||||
if verbose:
|
||||
print(f"[INFO] mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}")
|
||||
|
||||
return verts, faces
|
||||
|
||||
|
||||
def clean_mesh(
|
||||
verts,
|
||||
faces,
|
||||
v_pct=1,
|
||||
min_f=64,
|
||||
min_d=20,
|
||||
repair=True,
|
||||
remesh=True,
|
||||
remesh_size=0.01,
|
||||
):
|
||||
# verts: [N, 3]
|
||||
# faces: [N, 3]
|
||||
|
||||
_ori_vert_shape = verts.shape
|
||||
_ori_face_shape = faces.shape
|
||||
|
||||
m = pml.Mesh(verts, faces)
|
||||
ms = pml.MeshSet()
|
||||
ms.add_mesh(m, "mesh") # will copy!
|
||||
|
||||
# filters
|
||||
ms.meshing_remove_unreferenced_vertices() # verts not refed by any faces
|
||||
|
||||
if v_pct > 0:
|
||||
ms.meshing_merge_close_vertices(
|
||||
threshold=pml.Percentage(v_pct)
|
||||
) # 1/10000 of bounding box diagonal
|
||||
|
||||
ms.meshing_remove_duplicate_faces() # faces defined by the same verts
|
||||
ms.meshing_remove_null_faces() # faces with area == 0
|
||||
|
||||
if min_d > 0:
|
||||
ms.meshing_remove_connected_component_by_diameter(
|
||||
mincomponentdiag=pml.Percentage(min_d)
|
||||
)
|
||||
|
||||
if min_f > 0:
|
||||
ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f)
|
||||
|
||||
if repair:
|
||||
# ms.meshing_remove_t_vertices(method=0, threshold=40, repeat=True)
|
||||
ms.meshing_repair_non_manifold_edges(method=0)
|
||||
ms.meshing_repair_non_manifold_vertices(vertdispratio=0)
|
||||
|
||||
if remesh:
|
||||
# ms.apply_coord_taubin_smoothing()
|
||||
ms.meshing_isotropic_explicit_remeshing(
|
||||
iterations=3, targetlen=pml.AbsoluteValue(remesh_size)
|
||||
)
|
||||
|
||||
# extract mesh
|
||||
m = ms.current_mesh()
|
||||
verts = m.vertex_matrix()
|
||||
faces = m.face_matrix()
|
||||
|
||||
print(
|
||||
f"[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}"
|
||||
)
|
||||
|
||||
return verts, faces
|
||||
|
||||
def construct_list_of_gs_attributes(features_dc, features_rest, scaling, rotation):
|
||||
l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
|
||||
for i in range(features_dc.shape[1]*features_dc.shape[2]):
|
||||
l.append('f_dc_{}'.format(i))
|
||||
for i in range(features_rest.shape[1]*features_rest.shape[2]):
|
||||
l.append('f_rest_{}'.format(i))
|
||||
l.append('opacity')
|
||||
for i in range(scaling.shape[1]):
|
||||
l.append('scale_{}'.format(i))
|
||||
for i in range(rotation.shape[1]):
|
||||
l.append('rot_{}'.format(i))
|
||||
return l
|
||||
|
||||
def calculate_max_sh_degree_from_gs_ply(plydata):
|
||||
extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
|
||||
#assert len(extra_f_names)!=3*(max_sh_degree + 1) ** 2 - 3:
|
||||
max_sh_degree = int(((len(extra_f_names) + 3) / 3) ** 0.5 - 1)
|
||||
return max_sh_degree, extra_f_names
|
||||
|
||||
def write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, list_of_attributes):
|
||||
dtype_full = [(attribute, 'f4') for attribute in list_of_attributes]
|
||||
|
||||
elements = np.empty(xyz.shape[0], dtype=dtype_full)
|
||||
attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
|
||||
elements[:] = list(map(tuple, attributes))
|
||||
el = PlyElement.describe(elements, 'vertex')
|
||||
return PlyData([el])
|
||||
|
||||
def read_gs_ply(plydata):
|
||||
xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
|
||||
np.asarray(plydata.elements[0]["y"]),
|
||||
np.asarray(plydata.elements[0]["z"])), axis=1)
|
||||
opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
|
||||
|
||||
features_dc = np.zeros((xyz.shape[0], 3, 1))
|
||||
features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
|
||||
features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
|
||||
features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
|
||||
|
||||
max_sh_degree, extra_f_names = calculate_max_sh_degree_from_gs_ply(plydata)
|
||||
|
||||
features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
|
||||
for idx, attr_name in enumerate(extra_f_names):
|
||||
features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
|
||||
# Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
|
||||
features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))
|
||||
|
||||
scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
|
||||
scales = np.zeros((xyz.shape[0], len(scale_names)))
|
||||
for idx, attr_name in enumerate(scale_names):
|
||||
scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
|
||||
|
||||
rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
|
||||
rots = np.zeros((xyz.shape[0], len(rot_names)))
|
||||
for idx, attr_name in enumerate(rot_names):
|
||||
rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
|
||||
|
||||
return xyz, features_dc, features_extra, opacities, scales, rots
|
||||
|
||||
def ply_to_points_cloud(plydata):
|
||||
xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
|
||||
|
||||
features_dc = np.transpose(features_dc, (0, 2, 1)) # equivalent of torch.transpose(features_dc, 1, 2)
|
||||
features_extra = np.transpose(features_extra, (0, 2, 1))
|
||||
shs = np.concatenate((features_dc, features_extra), axis=1)
|
||||
normals = np.zeros_like(xyz)
|
||||
pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals)
|
||||
return pcd
|
||||
|
||||
|
||||
def get_target_axis_and_scale(axis_string, scale_value=1.0):
|
||||
"""
|
||||
Coordinate system inverts when:
|
||||
1. Any of the axis inverts
|
||||
2. Two of the axises switch
|
||||
|
||||
If coordinate system inverts twice in a row then it will not be inverted
|
||||
"""
|
||||
axis_names = ["x", "y", "z"]
|
||||
|
||||
target_axis, target_scale, coordinate_invert_count = [], [], 0
|
||||
axis_switch_count = 0
|
||||
for i in range(len(axis_names)):
|
||||
s = axis_string[i]
|
||||
if s[0] == "-":
|
||||
target_scale.append(-scale_value)
|
||||
coordinate_invert_count += 1
|
||||
else:
|
||||
target_scale.append(scale_value)
|
||||
|
||||
new_axis_i = axis_names.index(s[1])
|
||||
if new_axis_i != i:
|
||||
axis_switch_count += 1
|
||||
target_axis.append(new_axis_i)
|
||||
|
||||
if axis_switch_count == 2:
|
||||
coordinate_invert_count += 1
|
||||
|
||||
return target_axis, target_scale, coordinate_invert_count
|
||||
|
||||
def switch_vector_axis(vector3s, target_axis):
|
||||
"""
|
||||
Example:
|
||||
vector3s = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3)
|
||||
|
||||
target_axis = (2, 0, 1) # or [2, 0, 1]
|
||||
vector3s[:, [0, 1, 2]] = vector3s[:, target_axis]
|
||||
|
||||
# Result: tensor([[3, 1, 2], [1, 3, 2], [1, 2, 3]])
|
||||
"""
|
||||
vector3s[:, [0, 1, 2]] = vector3s[:, target_axis]
|
||||
return vector3s
|
||||
|
||||
def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_invert_count):
|
||||
"""
|
||||
Args:
|
||||
target_axis (array): shape (3)
|
||||
target_scale (array): shape (3)
|
||||
"""
|
||||
xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
|
||||
normals = np.zeros_like(xyz)
|
||||
features_dc_2d = features_dc.reshape(features_dc.shape[0], features_dc.shape[1]*features_dc.shape[2])
|
||||
features_extra_2d = features_extra.reshape(features_extra.shape[0], features_extra.shape[1]*features_extra.shape[2])
|
||||
|
||||
target_scale = torch.tensor(target_scale).float().cuda()
|
||||
xyz = switch_vector_axis(torch.tensor(xyz).float().cuda() * target_scale, target_axis).detach().cpu().numpy()
|
||||
scales = switch_vector_axis(torch.tensor(scales).float().cuda(), target_axis).detach().cpu().numpy()
|
||||
|
||||
# change rotation representation from quaternion (w, x, y, z) to axis angle vector (x, y, z) to make swich axis easier
|
||||
rots_axis_angle = quaternion_to_axis_angle(torch.tensor(rots).float().cuda())
|
||||
rots_axis_angle = switch_vector_axis(rots_axis_angle * target_scale, target_axis)
|
||||
"""
|
||||
Since axis–angle vector is composed of axis (unit vector/direction) and clockwise radians angle (vector magnitude),
|
||||
so in order to invert the sign of angle when coordinate system inverts, we also need to invert the direction of axis–angle vector
|
||||
"""
|
||||
if coordinate_invert_count % 2 != 0:
|
||||
rots_axis_angle = -rots_axis_angle
|
||||
rots = axis_angle_to_quaternion(rots_axis_angle).detach().cpu().numpy()
|
||||
|
||||
return write_gs_ply(xyz, normals, features_dc_2d, features_extra_2d, opacities, scales, rots, construct_list_of_gs_attributes(features_dc, features_extra, scales, rots))
|
||||
|
||||
def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=False):
|
||||
"""
|
||||
Args:
|
||||
target_axis (array): shape (3)
|
||||
target_scale (array): shape (3)
|
||||
"""
|
||||
target_scale = torch.tensor(target_scale).float().cuda()
|
||||
mesh.v = switch_vector_axis(mesh.v * target_scale, target_axis)
|
||||
mesh.vn = switch_vector_axis(mesh.vn * target_scale, target_axis)
|
||||
if flip_normal:
|
||||
mesh.vn *= -1
|
||||
return mesh
|
||||
|
||||
|
||||
def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4):
|
||||
from mcubes import marching_cubes
|
||||
from kiui.mesh_utils import clean_mesh, decimate_mesh
|
||||
|
||||
|
||||
sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32)
|
||||
|
||||
X = torch.linspace(-1, 1, grid_size).split(S)
|
||||
Y = torch.linspace(-1, 1, grid_size).split(S)
|
||||
Z = torch.linspace(-1, 1, grid_size).split(S)
|
||||
|
||||
for xi, xs in enumerate(X):
|
||||
for yi, ys in enumerate(Y):
|
||||
for zi, zs in enumerate(Z):
|
||||
xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij')
|
||||
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3]
|
||||
val = get_density_func(pts)
|
||||
sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z]
|
||||
|
||||
print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})')
|
||||
|
||||
vertices, triangles = marching_cubes(sigmas, density_thresh)
|
||||
vertices = vertices / (grid_size - 1.0) * 2 - 1
|
||||
|
||||
# clean
|
||||
vertices = vertices.astype(np.float32)
|
||||
triangles = triangles.astype(np.int32)
|
||||
vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
|
||||
if triangles.shape[0] > decimate_target:
|
||||
vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
|
||||
|
||||
return vertices, triangles
|
||||
|
||||
def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2, batch_size=640000, device="cuda", force_cuda_rast=False, glctx=None):
|
||||
import nvdiffrast.torch as dr
|
||||
from kiui.op import uv_padding
|
||||
|
||||
if glctx is None:
|
||||
if force_cuda_rast:
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
else:
|
||||
glctx = dr.RasterizeGLContext()
|
||||
|
||||
# render uv maps
|
||||
h = w = texture_resolution
|
||||
if mesh.vt is None:
|
||||
mesh.auto_uv()
|
||||
|
||||
uv = mesh.vt * 2.0 - 1.0 # uvs to range [-1, 1]
|
||||
uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4]
|
||||
|
||||
rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), mesh.ft, (h, w)) # [1, h, w, 4]
|
||||
xyzs, _ = dr.interpolate(mesh.v.unsqueeze(0), rast, mesh.f) # [1, h, w, 3]
|
||||
mask, _ = dr.interpolate(torch.ones_like(mesh.v[:, :1]).unsqueeze(0), rast, mesh.f) # [1, h, w, 1]
|
||||
|
||||
# masked query
|
||||
xyzs = xyzs.view(-1, 3)
|
||||
mask = (mask > 0).view(-1)
|
||||
|
||||
albedo = torch.zeros(h * w, 3, device=device, dtype=torch.float32)
|
||||
|
||||
if mask.any():
|
||||
print(f"[INFO] querying texture...")
|
||||
|
||||
xyzs = xyzs[mask] # [M, 3]
|
||||
|
||||
# batched inference to avoid OOM
|
||||
batch = []
|
||||
head = 0
|
||||
while head < xyzs.shape[0]:
|
||||
tail = min(head + batch_size, xyzs.shape[0])
|
||||
batch.append(get_rgb_func(xyzs[head:tail]).float())
|
||||
head += batch_size
|
||||
|
||||
albedo[mask] = torch.cat(batch, dim=0)
|
||||
|
||||
albedo = albedo.view(h, w, -1)
|
||||
mask = mask.view(h, w)
|
||||
albedo = uv_padding(albedo, mask, padding)
|
||||
|
||||
return albedo
|
||||
|
||||
@torch.no_grad()
|
||||
def K_nearest_neighbors_func(
|
||||
points: torch.Tensor,
|
||||
K: int,
|
||||
query: Optional[torch.Tensor] = None,
|
||||
return_dist=False,
|
||||
skip_first_index=True,
|
||||
):
|
||||
if not pytorch3d_capable:
|
||||
raise ImportError("pytorch3d is not installed, which is required for KNN")
|
||||
|
||||
# query/points: Tensor of shape (N, P1/P2, D) giving a batch of N point clouds, each containing up to P1/P2 points of dimension D
|
||||
if query is None:
|
||||
query = points
|
||||
dist, idx, nn = knn_points(query[None, ...], points[None, ...], K=K, return_nn=True)
|
||||
|
||||
# idx: Tensor of shape (N, P1, K)
|
||||
# nn: Tensor of shape (N, P1, K, D)
|
||||
|
||||
# take the index 1 since index 0 is the point itself
|
||||
if skip_first_index:
|
||||
if not return_dist:
|
||||
return nn[0, :, 1:, :], idx[0, :, 1:]
|
||||
else:
|
||||
return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:]
|
||||
else:
|
||||
if not return_dist:
|
||||
return nn[0, :, :, :], idx[0, :, :]
|
||||
else:
|
||||
return nn[0, :, :, :], idx[0, :, :], dist[0, :, :]
|
||||
|
||||
|
||||
def reduce_mesh_faces_inplace(mesh_obj, max_faces: int = 40000, backend="pymeshlab", optimalplacement=True, verbose=True):
|
||||
"""
|
||||
Reduce faces for Mesh/FastMesh objects using decimate_mesh (in-place modification)
|
||||
|
||||
Args:
|
||||
mesh_obj: Mesh or FastMesh object (modified in-place)
|
||||
max_faces: maximum number of faces
|
||||
backend: algorithm backend ("pymeshlab" or "pyfqmr")
|
||||
optimalplacement: optimal placement for decimation
|
||||
verbose: print progress
|
||||
|
||||
Returns:
|
||||
Same mesh object (modified in-place)
|
||||
"""
|
||||
|
||||
faces_before = mesh_obj.f.shape[0]
|
||||
|
||||
# Skip if already under limit
|
||||
if faces_before <= max_faces:
|
||||
if verbose:
|
||||
print(f"Mesh already has {faces_before} faces (≤ {max_faces}) - skipping")
|
||||
return mesh_obj
|
||||
|
||||
try:
|
||||
|
||||
# Convert to numpy
|
||||
vertices_np = mesh_obj.v.detach().cpu().numpy().astype(np.float32)
|
||||
faces_np = mesh_obj.f.detach().cpu().numpy().astype(np.int32)
|
||||
|
||||
# Apply decimation
|
||||
reduced_vertices, reduced_faces = decimate_mesh(
|
||||
vertices_np, faces_np,
|
||||
target=max_faces,
|
||||
backend=backend,
|
||||
optimalplacement=optimalplacement,
|
||||
verbose=verbose
|
||||
)
|
||||
|
||||
# Update existing mesh object (in-place)
|
||||
device = mesh_obj.v.device
|
||||
mesh_obj.v = torch.from_numpy(reduced_vertices).contiguous().float().to(device)
|
||||
mesh_obj.f = torch.from_numpy(reduced_faces).contiguous().int().to(device)
|
||||
|
||||
# Recalculate normals
|
||||
mesh_obj.auto_normal()
|
||||
|
||||
# Stats
|
||||
if verbose:
|
||||
faces_after = mesh_obj.f.shape[0]
|
||||
reduction_ratio = ((faces_before - faces_after) / faces_before) * 100
|
||||
|
||||
return mesh_obj
|
||||
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
print(f"Face reduction error: {e}, keep original mesh")
|
||||
return mesh_obj
|
||||
|
||||
|
||||
def interpolate_texture_map_attr(mesh, texture_size: int = 256, batch_size: int = 64, interpolate_color=True, interpolate_position=False):
|
||||
# Get UV coordinates and faces
|
||||
if mesh.vt is None:
|
||||
mesh.auto_uv()
|
||||
|
||||
# Get Faces on UV
|
||||
texture_size_minus_one = texture_size - 1
|
||||
faces = mesh.ft
|
||||
verts_uvs = mesh.vt * texture_size_minus_one
|
||||
verts_uvs_idx = verts_uvs.to(torch.long)
|
||||
|
||||
vmapping = mesh.get_default_vt_to_v_mapping()
|
||||
# Get vertex colors
|
||||
interpolate_color = interpolate_color and mesh.vc is not None
|
||||
texture_map = mesh.albedo
|
||||
if interpolate_color:
|
||||
verts_colors = mesh.vc[vmapping]
|
||||
# Create a blank texture map
|
||||
texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device)
|
||||
|
||||
# Get vertex positions
|
||||
position_map = None
|
||||
if interpolate_position:
|
||||
verts_positions = mesh.v[vmapping]
|
||||
position_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device)
|
||||
|
||||
# Create a grid of UV coordinates
|
||||
grid_x, grid_y = torch.meshgrid(torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device), torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device))
|
||||
grid = torch.stack([grid_x, grid_y], dim=-1)
|
||||
|
||||
batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long)
|
||||
|
||||
all_face_verts = verts_uvs_idx[faces]
|
||||
for uv_i in range(0, texture_size_minus_one, batch_size):
|
||||
for uv_j in range(0, texture_size_minus_one, batch_size):
|
||||
# Divide faces into squre batchs
|
||||
faces_mask = (uv_i <= all_face_verts[:, :, 0]) & (all_face_verts[:, :, 0] <= uv_i + batch_size) & (uv_j <= all_face_verts[:, :, 1]) & (all_face_verts[:, :, 1] <= uv_j + batch_size)
|
||||
faces_mask, _ = torch.max(faces_mask, dim=1)
|
||||
faces_mask = faces_mask.nonzero().squeeze(-1)
|
||||
batch_faces = faces[faces_mask]
|
||||
|
||||
uvs = verts_uvs[batch_faces]
|
||||
if interpolate_color:
|
||||
colors = verts_colors[batch_faces]
|
||||
if interpolate_position:
|
||||
positions = verts_positions[batch_faces]
|
||||
|
||||
v0, v1, v2 = uvs[:, 0], uvs[:, 1], uvs[:, 2]
|
||||
v0_0, v0_1 = v0[:, 0], v0[:, 1]
|
||||
v1_0, v1_1 = v1[:, 0], v1[:, 1]
|
||||
v2_0, v2_1 = v2[:, 0], v2[:, 1]
|
||||
|
||||
sub_grid = grid[uv_i:uv_i + batch_size, uv_j:uv_j + batch_size].reshape(-1, 2).unsqueeze(1)
|
||||
|
||||
# Compute barycentric coordinates for the batch of faces
|
||||
denom = (v1_1 - v2_1) * (v0_0 - v2_0) + (v2_0 - v1_0) * (v0_1 - v2_1)
|
||||
a = ((v1_1 - v2_1) * (sub_grid[:, :, 0] - v2_0) + (v2_0 - v1_0) * (sub_grid[:, :, 1] - v2_1)) / denom
|
||||
b = ((v2_1 - v0_1) * (sub_grid[:, :, 0] - v2_0) + (v0_0 - v2_0) * (sub_grid[:, :, 1] - v2_1)) / denom
|
||||
c = 1 - a - b
|
||||
|
||||
# Mask for points inside the triangle
|
||||
mask = (a >= 0) & (b >= 0) & (c >= 0)
|
||||
mask, mask_idx = torch.max(mask, dim=1)
|
||||
|
||||
# Get texture coordinates
|
||||
sub_grid = sub_grid.squeeze(1)
|
||||
uv_coords = sub_grid[mask].to(torch.long)
|
||||
|
||||
# Interpolate colors
|
||||
a, b, c = a[batch_idx_range, mask_idx][mask], b[batch_idx_range, mask_idx][mask], c[batch_idx_range, mask_idx][mask]
|
||||
|
||||
if interpolate_color:
|
||||
colors = colors[mask_idx][mask]
|
||||
interpolated_colors = (a[:, None] * colors[:, 0] + (b[:, None] * colors[:, 1]) + (c[:, None] * colors[:, 2]))
|
||||
texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors
|
||||
if interpolate_position:
|
||||
positions = positions[mask_idx][mask]
|
||||
interpolated_positions = (a[:, None] * positions[:, 0] + (b[:, None] * positions[:, 1]) + (c[:, None] * positions[:, 2]))
|
||||
position_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_positions
|
||||
|
||||
return texture_map, position_map
|
||||
@@ -37,8 +37,9 @@ from plyfile import PlyData
|
||||
import trimesh
|
||||
from PIL import Image
|
||||
|
||||
from .mesh_processer.mesh import Mesh
|
||||
from .mesh_processer.mesh_utils import (
|
||||
from mesh_processor.export_utils import export_to_fastmesh, export_to_mesh
|
||||
from mesh_processor.mesh import FastMesh, Mesh
|
||||
from mesh_processor.mesh_utils import (
|
||||
ply_to_points_cloud,
|
||||
get_target_axis_and_scale,
|
||||
switch_ply_axis_and_scale,
|
||||
@@ -48,6 +49,7 @@ from .mesh_processer.mesh_utils import (
|
||||
color_func_to_albedo,
|
||||
interpolate_texture_map_attr,
|
||||
decimate_mesh,
|
||||
reduce_mesh_faces_inplace
|
||||
)
|
||||
|
||||
from FlexiCubes.flexicubes_trainer import FlexiCubesTrainer
|
||||
@@ -128,6 +130,7 @@ from Gen_3D_Modules.PartCrafter.partcrafter_src.pipelines.pipeline_partcrafter i
|
||||
from Gen_3D_Modules.PartCrafter.partcrafter_src.utils.data_utils import get_colored_mesh_composition
|
||||
from Gen_3D_Modules.PartCrafter.partcrafter_src.utils.render_utils import explode_mesh
|
||||
import zipfile
|
||||
from typing import Union
|
||||
|
||||
|
||||
os.environ['SPCONV_ALGO'] = 'native'
|
||||
@@ -4569,7 +4572,6 @@ class Hunyuan3D_V2_ShapeGen_MV:
|
||||
|
||||
return (Mesh.load_trimesh(given_mesh=mesh),)
|
||||
|
||||
#--------------------------------
|
||||
class Load_StableGen_Trellis_Pipeline:
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
RETURN_TYPES = ("DIFFUSERS_PIPE",)
|
||||
@@ -4660,7 +4662,6 @@ class Load_StableGen_StableX_Pipeline:
|
||||
pipe = self.__class__._build_pipe(repo, use_fp16)
|
||||
return (pipe,)
|
||||
|
||||
|
||||
class StableGen_Trellis_Image_To_3D:
|
||||
"""
|
||||
3D generation pipeline using Trellis model.
|
||||
@@ -4766,7 +4767,6 @@ class StableGen_Trellis_Image_To_3D:
|
||||
except Exception as e:
|
||||
raise Exception(f"[StableGen_Trellis_Image_To_3D] 3D generation failed: {str(e)}")
|
||||
|
||||
|
||||
class StableGen_StableX_Process_Image:
|
||||
"""
|
||||
Image processing pipeline using StableX model.
|
||||
@@ -4814,7 +4814,7 @@ class StableGen_StableX_Process_Image:
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"[StableGen_StableX_Process_Image] Image processing failed: {str(e)}")
|
||||
# --- MV START ----
|
||||
|
||||
class Load_MVAdapter_IG2MV_Pipeline:
|
||||
"""Loader pipeline for MV-Adapter (Image to Multi-View)"""
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
@@ -4851,7 +4851,6 @@ class Load_MVAdapter_IG2MV_Pipeline:
|
||||
vae_model = None if vae_model == "None" else vae_model
|
||||
lora_model = None if not lora_model else lora_model
|
||||
|
||||
# Подготавливаем параметры для пайплайна
|
||||
pipeline_kwargs = {
|
||||
"base_model": base_model,
|
||||
"vae_model": vae_model,
|
||||
@@ -4870,6 +4869,7 @@ class Load_MVAdapter_IG2MV_Pipeline:
|
||||
print("MV-Adapter IG2MV pipeline loaded successfully")
|
||||
return (pipe,)
|
||||
|
||||
|
||||
class MVAdapter_IG2MV:
|
||||
"""Generate multi-view images from single image and 3D mesh"""
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
@@ -4931,6 +4931,7 @@ class MVAdapter_IG2MV:
|
||||
return_images = pils_to_torch_imgs(images, device=DEVICE_STR)
|
||||
return (return_images,)
|
||||
|
||||
|
||||
class Load_MVAdapter_TG2MV_Pipeline:
|
||||
"""Loader pipeline for MV-Adapter Text-Guided to Multi-View (TG2MV)"""
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
@@ -5044,6 +5045,7 @@ class MVAdapter_TG2MV:
|
||||
return_images = pils_to_torch_imgs(images, device=DEVICE_STR)
|
||||
return (return_images,)
|
||||
|
||||
|
||||
class Load_MVAdapter_Texture_Pipeline:
|
||||
"""Load texture projection pipeline for MVAdapter"""
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
@@ -5206,6 +5208,7 @@ class MVAdapter_Texture_Projection:
|
||||
os.remove(temp_grid_path)
|
||||
raise e
|
||||
|
||||
|
||||
class Load_Hunyuan3D_21_ShapeGen_Pipeline:
|
||||
"""Load Hunyuan3D-2.1 Shape Generation Pipeline"""
|
||||
|
||||
@@ -5274,6 +5277,7 @@ class Load_Hunyuan3D_21_ShapeGen_Pipeline:
|
||||
|
||||
return (pipeline,)
|
||||
|
||||
|
||||
class Load_Hunyuan3D_21_TexGen_Pipeline:
|
||||
"""Load Hunyuan3D-2.1 Texture Generation Pipeline"""
|
||||
|
||||
@@ -5394,8 +5398,9 @@ class Load_Hunyuan3D_21_TexGen_Pipeline:
|
||||
print(f"[TexGen-Loader] Cached new pipeline {cache_key}")
|
||||
return (pipeline,)
|
||||
|
||||
|
||||
class Hunyuan3D_21_ShapeGen:
|
||||
"""Hunyuan3D-2.1 Shape Generation with automatic pipeline cleanup"""
|
||||
"""Hunyuan3D-2.1 Shape Generation """
|
||||
|
||||
CATEGORY = "Comfy3D/Algorithm/Hunyuan3D-2.1"
|
||||
RETURN_TYPES = ("MESH", "IMAGE")
|
||||
@@ -5414,20 +5419,20 @@ class Hunyuan3D_21_ShapeGen:
|
||||
"octree_resolution": ("INT", {"default": 256, "min": 64, "max": 512}),
|
||||
"remove_background": ("BOOLEAN", {"default": True}),
|
||||
"auto_cleanup": ("BOOLEAN", {"default": True}),
|
||||
"max_faces": ("INT", {"default": 40000, "min": 1000, "max": 1000000}),
|
||||
}
|
||||
}
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, shapegen_pipe, image, seed, steps, guidance_scale, octree_resolution, remove_background, auto_cleanup):
|
||||
pil_image = torch_imgs_to_pils(image)[0].convert("RGBA")
|
||||
|
||||
if remove_background or pil_image.mode == "RGB":
|
||||
rmbg_worker = BackgroundRemover_2_1()
|
||||
pil_image = rmbg_worker(pil_image.convert('RGB'))
|
||||
del rmbg_worker
|
||||
|
||||
generator = torch.Generator(device=shapegen_pipe.device)
|
||||
generator = generator.manual_seed(int(seed))
|
||||
pil_image = torch_imgs_to_pils(image)[0]
|
||||
if remove_background:
|
||||
bg_remover = BackgroundRemover_2_1()
|
||||
if bg_remover is not None:
|
||||
pil_image = bg_remover(pil_image)
|
||||
del bg_remover
|
||||
|
||||
generator = torch.Generator(device=shapegen_pipe.device).manual_seed(seed)
|
||||
|
||||
outputs = shapegen_pipe(
|
||||
image=pil_image,
|
||||
@@ -5439,36 +5444,23 @@ class Hunyuan3D_21_ShapeGen:
|
||||
output_type='mesh'
|
||||
)
|
||||
|
||||
mesh = export_to_trimesh_2_1(outputs)[0]
|
||||
mesh_objects = export_to_fastmesh(outputs)
|
||||
mesh_out = mesh_objects[0] if isinstance(mesh_objects, list) else mesh_objects
|
||||
|
||||
face_reduce_worker = FaceReducer_2_1()
|
||||
mesh = face_reduce_worker(mesh)
|
||||
del face_reduce_worker
|
||||
if mesh_out is None:
|
||||
raise Exception("Failed to create mesh from pipeline")
|
||||
|
||||
reduce_mesh_faces_inplace(mesh_out, max_faces=max_faces, verbose=True)
|
||||
|
||||
# Auto cleanup pipeline if enabled
|
||||
if auto_cleanup:
|
||||
try:
|
||||
shapegen_pipe.to('cpu')
|
||||
if hasattr(shapegen_pipe, 'unet'):
|
||||
del shapegen_pipe.unet
|
||||
if hasattr(shapegen_pipe, 'vae'):
|
||||
del shapegen_pipe.vae
|
||||
if hasattr(shapegen_pipe, 'scheduler'):
|
||||
del shapegen_pipe.scheduler
|
||||
del outputs
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
print("Shape pipeline cleaned up")
|
||||
except Exception as e:
|
||||
print(f"Error during pipeline cleanup: {e}")
|
||||
|
||||
mesh_out = Mesh.load_trimesh(given_mesh=mesh)
|
||||
mesh_out.auto_normal()
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
processed_image_tensor = pils_to_torch_imgs([pil_image])
|
||||
|
||||
return (mesh_out, processed_image_tensor)
|
||||
|
||||
|
||||
class Hunyuan3D_21_TexGen:
|
||||
"""Hunyuan3D-2.1 Texture Generation"""
|
||||
|
||||
@@ -5587,7 +5579,6 @@ class Hunyuan3D_21_TexGen:
|
||||
except:
|
||||
pass
|
||||
|
||||
# --------------------- PARTCRAFTER ---------------------
|
||||
|
||||
class Load_PartCrafter_Pipeline:
|
||||
"""Load PartCrafter Pipeline"""
|
||||
@@ -5790,7 +5781,7 @@ class PartCrafter_Generate:
|
||||
print(f"GLB mesh path for Preview_3DMesh: {relative_scene_path}")
|
||||
|
||||
return (zip_path, relative_scene_path, processed_image_tensor)
|
||||
#------ partcrafter scene ---------------------
|
||||
|
||||
|
||||
class Load_PartCrafter_Scene_Pipeline:
|
||||
"""Load PartCrafter Scene Pipeline"""
|
||||
@@ -5994,4 +5985,3 @@ class PartCrafter_Generate:
|
||||
print(f"GLB mesh path for Preview_3DMesh: {relative_scene_path}")
|
||||
|
||||
return (zip_path, relative_scene_path, processed_image_tensor)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user