nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Finished GPU accelerated Convert_Vertex_Color_To_Texture function

Browse Source
This commit is contained in:
Mr. For Example
2024-08-21 16:53:15 -06:00
parent 945fd06f6d
commit df3b12e3c4
4 changed files with 48 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
_Pre_Builds/_Build_Scripts/auto_build_all.py
+1 -1
View File
@@ -45,7 +45,7 @@ def get_dependency_dir(dependency):
def setup_build_env():
# Set CMake argumens fo build packages based on CUDA
os.environ["CMAKE_ARGS"]="-DBUILD_opencv_world=ON -DWITH_CUDA=ON -DCUDA_FAST_MATH=ON -DWITH_CUBLAS=ON -DCUDA_ARCH_PTX=9.0 -DWITH_NVCUVID=ON"
os.environ["CMAKE_ARGS"] = "-DBUILD_opencv_world=ON -DWITH_CUDA=ON -DCUDA_FAST_MATH=ON -DWITH_CUBLAS=ON -DCUDA_ARCH_PTX=9.0 -DWITH_NVCUVID=ON"
subprocess.run([PYTHON_PATH, "-s", "-m", "pip", "install", "-r", BUILD_REQUIREMENTS_FILE_ABS_PATH])
install_remote_packages(build_config.build_base_packages)
mesh_processer/mesh.py
+15 -4
View File
@@ -540,10 +540,7 @@ class Mesh:
vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None.
"""
if vmapping is None:
ft = self.ft.view(-1).long()
f = self.f.view(-1).long()
vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device)
vmapping[ft] = f # scatter, randomly choose one if index is not unique
vmapping = self.get_default_vmapping()
self.v = self.v[vmapping]
if self.vc is not None:
@@ -553,6 +550,20 @@ class Mesh:
if self.vn is not None:
self.vn = self.vn[vmapping]
self.fn = self.ft
print(f"########## !!!!!!!!!! vmapping: {vmapping.shape}; self.vt.shape: {self.vt.shape}; self.v: {self.v.shape}; self.f: {self.f.shape}")
def get_default_vmapping(self):
"""map from ft to f
Returns:
vmapping: tensor array
"""
ft = self.ft.view(-1).long()
f = self.f.view(-1).long()
vmapping = torch.zeros(self.vt.shape[0], dtype=torch.long, device=self.device)
vmapping[ft] = f # scatter, randomly choose one if index is not unique
return vmapping
def to(self, device):
"""move all tensor attributes to device.
mesh_processer/mesh_utils.py
+20 -99
View File
@@ -582,115 +582,35 @@ def K_nearest_neighbors_func(
return nn[0, :, :, :], idx[0, :, :]
else:
return nn[0, :, :, :], idx[0, :, :], dist[0, :, :]
def _interpolate_texture_map(mesh, texture_size: int = 256):
# Get UV coordinates and faces
if mesh.vt is None:
mesh.auto_uv()
verts_uvs = mesh.vt
verts_uvs_idx = (verts_uvs * (texture_size - 1)).to(torch.int)
faces = mesh.f
verts_colors = mesh.vc
# Create a blank texture map
texture_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, 1, texture_size, device=verts_colors.device), torch.linspace(0, 1, texture_size, device=verts_colors.device))
#grid = torch.stack([grid_x, grid_y], dim=-1).reshape(-1, 2)
grid = torch.stack([grid_x, grid_y], dim=-1)
# Function to interpolate colors within a triangle
def _interpolate_triangle(grid, uvs, colors):
# Compute barycentric coordinates
v0, v1, v2 = uvs
denom = (v1[1] - v2[1]) * (v0[0] - v2[0]) + (v2[0] - v1[0]) * (v0[1] - v2[1])
a = ((v1[1] - v2[1]) * (grid[:, 0] - v2[0]) + (v2[0] - v1[0]) * (grid[:, 1] - v2[1])) / denom
b = ((v2[1] - v0[1]) * (grid[:, 0] - v2[0]) + (v0[0] - v2[0]) * (grid[:, 1] - v2[1])) / denom
c = 1 - a - b
# Mask for points inside the triangle
mask = (a >= 0) & (b >= 0) & (c >= 0)
uv_coords = (grid[mask] * (texture_size - 1)).to(torch.int)
# Interpolate colors
interpolated_colors = a[:, None] * colors[0] + b[:, None] * colors[1] + c[:, None] * colors[2]
#print(f"############sub_grid: {grid.shape}; a: {a.shape}; b: {b.shape}; mask: {mask.shape}; uv_coords: {grid[mask].shape}; interpolated_colors: {interpolated_colors[mask].shape}")
#return grid[mask], interpolated_colors[mask]
return uv_coords, interpolated_colors[mask]
# Function to interpolate colors within a triangle
def interpolate_triangle(grid, uvs_idx, uvs, colors):
u_idx, v_idx = uvs_idx[:, 0], uvs_idx[:, 1]
sub_grid = grid[torch.min(u_idx):torch.max(u_idx)+1, torch.min(v_idx):torch.max(v_idx)+1]
#sub_grid = grid
# Compute barycentric coordinates
v0, v1, v2 = uvs
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)
# Get texture coordinates
#uv_coords = (sub_grid[mask][[1,0]] * (texture_size - 1)).to(torch.int)
uv_coords = (sub_grid[mask] * (texture_size - 1)).to(torch.int)
# Interpolate colors
interpolated_colors = (a[:, :, None] * colors[0] + b[:, :, None] * colors[1] + c[:, :, None] * colors[2])[mask]
#print(f"############sub_grid: {sub_grid.shape}; a: {a.shape}; b: {b.shape}; mask: {mask.shape}; uv_coords: {uv_coords.shape}; interpolated_colors: {interpolated_colors.shape}")
return uv_coords, interpolated_colors
# Map vertex colors to the UV coordinates
import comfy.utils
comfy_pbar = comfy.utils.ProgressBar(faces.shape[0])
for i, face in enumerate(faces):
uvs = verts_uvs[face]
uvs_idx = verts_uvs_idx[face]
colors = verts_colors[face]
uv_coords, interpolated_colors = interpolate_triangle(grid, uvs_idx, uvs, colors)
#uv_coords, interpolated_colors = interpolate_triangle(grid, uvs, colors)
# Map interpolated colors to the texture map
#for uv, color in zip(uv_coords, interpolated_colors):
# x, y = int(uv[0] * (texture_size - 1)), int(uv[1] * (texture_size - 1))
# texture_map[y, x] = color
texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors
comfy_pbar.update_absolute(i + 1)
return texture_map
def interpolate_texture_map(mesh, texture_size: int = 256, batch_size: int = 64):
# Get UV coordinates and faces
if mesh.vt is None:
mesh.auto_uv()
verts_uvs = mesh.vt
verts_uvs_idx = (mesh.vt * (texture_size - 1)).to(torch.int)
faces = mesh.f
verts_colors = mesh.vc
# 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)
# Get vertex colors
vmapping = mesh.get_default_vmapping()
verts_colors = mesh.vc[vmapping]
# Create a blank texture map
texture_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, 1, texture_size, device=verts_colors.device), torch.linspace(0, 1, texture_size, device=verts_colors.device))
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.int)
batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long)
# Divide faces into squre batchs
all_face_verts = verts_uvs_idx[faces]
for uv_i in range(0, texture_size-1, batch_size):
for uv_j in range(0, texture_size-1, batch_size):
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)
@@ -698,14 +618,15 @@ def interpolate_texture_map(mesh, texture_size: int = 256, batch_size: int = 64)
uvs = verts_uvs[batch_faces]
colors = verts_colors[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
@@ -714,10 +635,10 @@ def interpolate_texture_map(mesh, texture_size: int = 256, batch_size: int = 64)
# 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] * (texture_size - 1)).to(torch.int)
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]
nodes.py
+12 -18
View File
@@ -43,7 +43,6 @@ from .mesh_processer.mesh_utils import (
calculate_max_sh_degree_from_gs_ply,
marching_cubes_density_to_mesh,
color_func_to_albedo,
K_nearest_neighbors_func,
interpolate_texture_map,
)
@@ -3018,10 +3017,15 @@ class ExplicitTarget_Color_Projection:
new_meshes = multiview_color_projection(meshes, pil_image_list, weights=weights, resolution=projection_resolution, device=DEVICE, complete_unseen=complete_unseen_rgb, confidence_threshold=confidence_threshold, cameras_list=cam_list)
vertices, faces, vertex_colors = from_py3d_mesh(new_meshes)
vertices = vertices / 2 * 1.35
mesh = Mesh(v=vertices, f=faces, vc=vertex_colors, device=DEVICE)
mesh.auto_normal()
print(f"###############faces: {faces.shape} \nmesh.f: {mesh.f.shape} \nvertices: {vertices.shape} \nmesh.v: {mesh.v.shape} \nmesh.vt: {mesh.vt.shape}")
mesh = Mesh(v=vertices, f=faces,
vn=None if mesh.vn is None else mesh.vn.clone(), fn=None if mesh.fn is None else mesh.fn.clone(),
vt=None if mesh.vt is None else mesh.vt.clone(), ft=None if mesh.ft is None else mesh.ft.clone(),
vc=vertex_colors, device=DEVICE)
if mesh.vn is None:
mesh.auto_normal()
return (mesh,)
class Convert_Vertex_Color_To_Texture:
@@ -3031,8 +3035,7 @@ class Convert_Vertex_Color_To_Texture:
"required": {
"mesh": ("MESH",),
"texture_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"K_nearest_neighbors": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"force_cuda_rast": ("BOOLEAN", {"default": False}),
"batch_size": ("INT", {"default": 128, "min": 1, "max": 0xffffffffffffffff}),
},
}
@@ -3045,25 +3048,16 @@ class Convert_Vertex_Color_To_Texture:
FUNCTION = "run_convert_func"
CATEGORY = "Comfy3D/Algorithm"
def run_convert_func(self, mesh, texture_resolution, K_nearest_neighbors, force_cuda_rast):
def run_convert_func(self, mesh, texture_resolution, batch_size):
if mesh.vc is not None:
#self.v = mesh.v.detach()
#self.vc = mesh.vc.detach()
#self.K = K_nearest_neighbors
#mesh.albedo = color_func_to_albedo(mesh, self.get_color_func, texture_resolution, device=DEVICE, force_cuda_rast=force_cuda_rast)
mesh.albedo = interpolate_texture_map(mesh, texture_resolution)
albedo_img = interpolate_texture_map(mesh, texture_resolution, batch_size)
mesh.albedo = troch_image_dilate(albedo_img)
else:
cstr(f"[{self.__class__.__name__}] skip this node since there is no vertex color found in mesh").msg.print()
return (mesh,)
def get_color_func(self, xs):
_, idx, dist = K_nearest_neighbors_func(self.v, self.K, xs, return_dist=True, skip_first_index=False)
colors = torch.sum(self.vc[idx] * (dist / torch.sum(dist, dim=-1, keepdim=True)).unsqueeze(-1), dim=1)
return colors
class Load_CharacterGen_MVDiffusion_Model:
checkpoints_dir = "CharacterGen"
default_repo_id = "zjpshadow/CharacterGen"