Add three draft of Convert_Vertex_Color_To_Texture function

This commit is contained in:
Mr. For Example
2024-08-21 12:27:07 -06:00
parent 17584d51a0
commit 945fd06f6d
8 changed files with 293 additions and 75 deletions
@@ -91,7 +91,6 @@ class Pix2FacesRenderer:
cameras = cameras.to(self.device)
vertices = self.transform_vertices(meshes, cameras)
faces = meshes.faces_packed().to(torch.int32)
print(f"#############vertices: {vertices.shape}; faces: {faces.shape}")
rast_out,_ = dr.rasterize(self._glctx, vertices, faces, resolution=(H, W), grad_db=False) #C,H,W,4
pix_to_face = rast_out[..., -1].to(torch.int32) - 1
return pix_to_face
@@ -267,12 +266,10 @@ def multiview_color_projection(meshes: Meshes, image_list: List[Image.Image], ca
elif len(image_list) == 2:
weights = [1.0, 1.0]
else:
raise ValueError("weights is None, and can not be guessed from image_list")
weights = [1. for _ in range(len(cameras_list))]
# 2. run projection
meshes = meshes.clone().to(device)
if weights is None:
weights = [1. for _ in range(len(cameras_list))]
assert len(cameras_list) == len(image_list) == len(weights)
original_color = meshes.textures.verts_features_packed()
assert not torch.isnan(original_color).any()
@@ -315,10 +312,10 @@ def get_orbit_cameras_list(orbit_camera_poses, device, fov_in_degrees=60):
ret = []
for i in range(len(orbit_camera_poses)):
campose = orbit_camera_poses[i]
R, T = look_at_view_transform(campose[0], campose[1], campose[2])
R, T = look_at_view_transform(campose[0], -campose[1], campose[2])
w2c = torch.cat([R[0].T, T[0, :, None]], dim=1)
cameras: PerspectiveCameras = get_camera(w2c, fov_in_degrees=fov_in_degrees, cam_type='fov').to(device)
ret.append(cameras)
camera: PerspectiveCameras = get_camera(w2c, fov_in_degrees=fov_in_degrees, cam_type='fov').to(device)
ret.append(camera)
return ret
def get_cameras_list(azim_list, device, focal=2/1.35, dist=1.1):
@@ -326,8 +323,8 @@ def get_cameras_list(azim_list, device, focal=2/1.35, dist=1.1):
for azim in azim_list:
R, T = look_at_view_transform(dist, 0, azim)
w2c = torch.cat([R[0].T, T[0, :, None]], dim=1)
cameras: OrthographicCameras = get_camera(w2c, focal_length=focal, cam_type='orthogonal').to(device)
ret.append(cameras)
camera: OrthographicCameras = get_camera(w2c, focal_length=focal, cam_type='orthogonal').to(device)
ret.append(camera)
return ret
def get_8view_cameras(device, focal=2/1.35):
+9 -5
View File
@@ -81,14 +81,18 @@ def to_pyml_mesh(vertices,faces):
return m1
def to_py3d_mesh(vertices, faces, normals=None):
def to_py3d_mesh(vertices, faces, normals=None, vertices_color=None):
from pytorch3d.structures import Meshes
from pytorch3d.renderer.mesh.textures import TexturesVertex
mesh = Meshes(verts=[vertices], faces=[faces], textures=None)
if normals is None:
normals = mesh.verts_normals_packed()
# set normals as vertext colors
mesh.textures = TexturesVertex(verts_features=[normals / 2 + 0.5])
if vertices_color is None:
# set normals as vertext colors
if normals is None:
normals = mesh.verts_normals_packed()
mesh.textures = TexturesVertex(verts_features=[normals / 2 + 0.5])
else:
mesh.textures = TexturesVertex(verts_features=[vertices_color])
return mesh
@@ -1,21 +1,12 @@
import os
import math
import random
import numpy as np
from typing import Optional
from torchtyping import TensorType
from plyfile import PlyData, PlyElement
import torch
from torch import nn
import torch.nn.functional as F
pytorch3d_capable = True
try:
import pytorch3d
from pytorch3d.ops import knn_points
except ImportError:
pytorch3d_capable = False
from kornia.geometry.conversions import (
quaternion_to_rotation_matrix,
@@ -31,7 +22,7 @@ from kiui.op import inverse_sigmoid
from shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
from mesh_processer.mesh import Mesh, PointCloud
from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply
from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply, K_nearest_neighbors_func
def get_expon_lr_func(
lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
@@ -170,31 +161,6 @@ def find_points_within_radius(query_points, vertex_points, d):
return result
@torch.no_grad()
def K_nearest_neighbors(
points: torch.Tensor,
K: int,
query: Optional[torch.Tensor] = None,
return_dist=False,
):
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 not return_dist:
return nn[0, :, 1:, :], idx[0, :, 1:]
else:
return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:]
def distance_to_gaussian_surface(points, svec, rotmat, query):
"""
Calculate the radius of the gaussian along the direction, which determined by offset between points and query
@@ -773,7 +739,7 @@ class GaussianModel:
return new_params
def densify_by_compatness(self, K=1):
_, idx = K_nearest_neighbors(self._xyz, K=K+1)
_, idx = K_nearest_neighbors_func(self._xyz, K=K+1)
new_params_list = []
for i in range(K):
+3 -3
View File
@@ -17,11 +17,11 @@ build_base_packages: ["torch", "torchvision"]
# All packages will be automatically installed at beginning of install.py follwing orders from top to bottom
remote_packages:
torch:
version: "2.3.0"
version: "2.4.0"
url: "https://download.pytorch.org/whl/${cuda_version}"
torchvision:
version: "0.18.0"
version: "0.19.0"
url: "https://download.pytorch.org/whl/${cuda_version}"
torch-scatter:
@@ -29,5 +29,5 @@ remote_packages:
url: "https://data.pyg.org/whl/torch-${remote_packages.torch.version}+${cuda_version}.html"
xformers:
version: "0.0.26.post1"
version: "0.0.27.post2"
url: "https://download.pytorch.org/whl/${cuda_version}"
+2
View File
@@ -546,6 +546,8 @@ class Mesh:
vmapping[ft] = f # scatter, randomly choose one if index is not unique
self.v = self.v[vmapping]
if self.vc is not None:
self.vc = self.vc[vmapping]
self.f = self.ft
# assume fn == f
if self.vn is not None:
+188 -1
View File
@@ -7,6 +7,15 @@ from kornia.geometry.conversions import (
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
@@ -506,6 +515,9 @@ def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2,
# 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]
@@ -538,4 +550,179 @@ def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2,
mask = mask.view(h, w)
albedo = uv_padding(albedo, mask, padding)
return albedo
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 _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
# 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)
batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.int)
# 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):
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]
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)
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] * (texture_size - 1)).to(torch.int)
# 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]
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
return texture_map
+60 -19
View File
@@ -43,6 +43,8 @@ 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,
)
from FlexiCubes.flexicubes_trainer import FlexiCubesTrainer
@@ -90,7 +92,10 @@ from .shared_utils.image_utils import (
prepare_torch_img, torch_imgs_to_pils, troch_image_dilate,
pils_rgba_to_rgb, pil_make_image_grid, pil_split_image, pils_to_torch_imgs, pils_resize_foreground
)
from .shared_utils.camera_utils import compose_orbit_camposes
from .shared_utils.camera_utils import (
ORBITPOSE_PRESET_DICT, ELEVATION_MIN, ELEVATION_MAX, AZIMUTH_MIN, AZIMUTH_MAX,
compose_orbit_camposes
)
from .shared_utils.log_utils import cstr
from .shared_utils.common_utils import parse_save_filename, get_list_filenames, resume_or_download_model_from_hf
@@ -116,18 +121,6 @@ DIFFUSERS_SCHEDULER_DICT = OrderedDict([
("KDPM2DiscreteScheduler,", KDPM2DiscreteScheduler),
])
#{Key: [elevation, azimuth], ...}
ORBITPOSE_PRESET_DICT = OrderedDict([
("Custom", [[0.0, 90.0, 0.0, 0.0, -90.0, 0.0], [-90.0, 0.0, 180.0, 90.0, 0.0, 0.0]]),
("CRM(6)", [[0.0, 90.0, 0.0, 0.0, -90.0, 0.0], [-90.0, 0.0, 180.0, 90.0, 0.0, 0.0]]),
("Wonder3D(6)", [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 45.0, 90.0, 180.0, -90.0, -45.0]]),
("Zero123Plus(6)", [[-20.0, 10.0, -20.0, 10.0, -20.0, 10.0], [30.0, 90.0, 150.0, -150.0, -90.0, -30.0]]),
("Era3D(6)", [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 45.0, 90.0, 180.0, -90.0, -45.0]]),
("MVDream(4)", [[0.0, 0.0, 0.0, 0.0], [0.0, 90.0, 180.0, -90.0]]),
("Unique3D(4)", [[0.0, 0.0, 0.0, 0.0], [0.0, 90.0, 180.0, -90.0]]),
("CharacterGen(4)", [[0.0, 0.0, 0.0, 0.0], [-90.0, 180.0, 90.0, 0.0]]),
])
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
CKPT_ROOT_PATH = os.path.join(ROOT_PATH, "Checkpoints")
CKPT_DIFFUSERS_PATH = os.path.join(CKPT_ROOT_PATH, "Diffusers")
@@ -158,11 +151,6 @@ SUPPORTED_CHECKPOINTS_EXTENSIONS = (
'.safetensors',
)
ELEVATION_MIN = -90
ELEVATION_MAX = 90.0
AZIMUTH_MIN = -180.0
AZIMUTH_MAX = 180.0
WEIGHT_DTYPE = torch.float16
DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
@@ -1297,6 +1285,9 @@ class Fitting_Mesh_With_Multiview_Images:
force_cuda_rasterize,
):
if mesh.vt is None:
mesh.auto_uv()
mesh.set_new_albedo(mesh_albedo_width, mesh_albedo_height)
trained_mesh = None
@@ -2971,6 +2962,8 @@ class ExplicitTarget_Color_Projection:
"projection_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"complete_unseen_rgb": ("BOOLEAN", {"default": True},),
"render_orbit_camera_fovy": ("FLOAT", {"default": 47.5, "min": 0.0, "max": 180.0, "step": 0.1}),
"projection_weights": ("STRING", {"default": "2.0, 0.2, 1.0, 0.2"}),
"confidence_threshold": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 1.0, "step": 0.001}),
},
"optional": {
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
@@ -2994,6 +2987,8 @@ class ExplicitTarget_Color_Projection:
projection_resolution,
complete_unseen_rgb,
render_orbit_camera_fovy,
projection_weights,
confidence_threshold,
reference_orbit_camera_poses=None,
):
pil_image_list = torch_imgs_to_pils(reference_images, reference_masks)
@@ -3015,7 +3010,13 @@ class ExplicitTarget_Color_Projection:
#reference_orbit_camera_poses[0] = [360 + angle if angle < 0 else angle for angle in reference_orbit_camera_poses[0]]
cam_list = get_orbit_cameras_list(reference_orbit_camera_poses, DEVICE, render_orbit_camera_fovy)
new_meshes = multiview_color_projection(meshes, pil_image_list, resolution=projection_resolution, device=DEVICE, complete_unseen=complete_unseen_rgb, confidence_threshold=0.2, cameras_list=cam_list)
weights = projection_weights.split(",")
if len(weights) == len(cam_list):
weights = [float(item) for item in weights]
else:
weights = None
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
@@ -3023,6 +3024,46 @@ class ExplicitTarget_Color_Projection:
mesh.auto_normal()
return (mesh,)
class Convert_Vertex_Color_To_Texture:
@classmethod
def INPUT_TYPES(cls):
return {
"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}),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_convert_func"
CATEGORY = "Comfy3D/Algorithm"
def run_convert_func(self, mesh, texture_resolution, K_nearest_neighbors, force_cuda_rast):
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)
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"
+23 -2
View File
@@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from collections import OrderedDict
import numpy as np
import math
from scipy.spatial.transform import Rotation as R
@@ -8,6 +8,22 @@ import torch
from kiui.cam import orbit_camera
#{Key: [elevation, azimuth], ...}
ORBITPOSE_PRESET_DICT = OrderedDict([
("Custom", [[0.0, 90.0, 0.0, 0.0, -90.0, 0.0], [-90.0, 0.0, 180.0, 90.0, 0.0, 0.0]]),
("CRM(6)", [[0.0, 90.0, 0.0, 0.0, -90.0, 0.0], [-90.0, 0.0, 180.0, 90.0, 0.0, 0.0]]),
("Wonder3D(6)", [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 45.0, 90.0, 180.0, -90.0, -45.0]]),
("Zero123Plus(6)", [[-20.0, 10.0, -20.0, 10.0, -20.0, 10.0], [30.0, 90.0, 150.0, -150.0, -90.0, -30.0]]),
("Era3D(6)", [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 45.0, 90.0, 180.0, -90.0, -45.0]]),
("MVDream(4)", [[0.0, 0.0, 0.0, 0.0], [0.0, 90.0, 180.0, -90.0]]),
("Unique3D(4)", [[0.0, 0.0, 0.0, 0.0], [0.0, 90.0, 180.0, -90.0]]),
("CharacterGen(4)", [[0.0, 0.0, 0.0, 0.0], [-90.0, 180.0, 90.0, 0.0]]),
])
ELEVATION_MIN = -89.999
ELEVATION_MAX = 89.999
AZIMUTH_MIN = -180.0
AZIMUTH_MAX = 180.0
def dot(x, y):
if isinstance(x, np.ndarray):
return np.sum(x * y, -1, keepdims=True)
@@ -262,6 +278,11 @@ def compose_orbit_camposes(orbit_radius, orbit_elevations, orbit_azimuths, orbit
campose_num = len(orbit_radius)
for i in range(campose_num):
orbit_camposes.append([orbit_radius[i], orbit_elevations[i], orbit_azimuths[i], orbit_center_x[i], orbit_center_y[i], orbit_center_z[i]])
orbit_camposes.append([
orbit_radius[i],
np.clip(orbit_elevations[i], ELEVATION_MIN, ELEVATION_MAX),
np.clip(orbit_azimuths[i], AZIMUTH_MIN, AZIMUTH_MAX),
orbit_center_x[i], orbit_center_y[i], orbit_center_z[i]
])
return orbit_camposes