feat: init git
This commit is contained in:
陈浩杰
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,258 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import numpy as np
|
||||
import math
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
import torch
|
||||
|
||||
from kiui.cam import orbit_camera
|
||||
|
||||
def dot(x, y):
|
||||
if isinstance(x, np.ndarray):
|
||||
return np.sum(x * y, -1, keepdims=True)
|
||||
else:
|
||||
return torch.sum(x * y, -1, keepdim=True)
|
||||
|
||||
|
||||
def length(x, eps=1e-20):
|
||||
if isinstance(x, np.ndarray):
|
||||
return np.sqrt(np.maximum(np.sum(x * x, axis=-1, keepdims=True), eps))
|
||||
else:
|
||||
return torch.sqrt(torch.clamp(dot(x, x), min=eps))
|
||||
|
||||
|
||||
def safe_normalize(x, eps=1e-20):
|
||||
return x / length(x, eps)
|
||||
|
||||
|
||||
def look_at(campos, target, opengl=True):
|
||||
# campos: [N, 3], camera/eye position
|
||||
# target: [N, 3], object to look at
|
||||
# return: [N, 3, 3], rotation matrix
|
||||
if not opengl:
|
||||
# camera forward aligns with -z
|
||||
forward_vector = safe_normalize(target - campos)
|
||||
up_vector = np.array([0, 1, 0], dtype=np.float32)
|
||||
right_vector = safe_normalize(np.cross(forward_vector, up_vector))
|
||||
up_vector = safe_normalize(np.cross(right_vector, forward_vector))
|
||||
else:
|
||||
# camera forward aligns with +z
|
||||
forward_vector = safe_normalize(campos - target)
|
||||
up_vector = np.array([0, 1, 0], dtype=np.float32)
|
||||
right_vector = safe_normalize(np.cross(up_vector, forward_vector))
|
||||
up_vector = safe_normalize(np.cross(forward_vector, right_vector))
|
||||
R = np.stack([right_vector, up_vector, forward_vector], axis=1)
|
||||
return R
|
||||
|
||||
def get_look_at_camera_pose(target, target_to_cam_offset, look_distance=0.1, opengl=True):
|
||||
"""
|
||||
Calculate the pose (cam2world) matrix from target position the camera suppose to look at and offset vector from target to camera
|
||||
|
||||
Args:
|
||||
target (NDArray[float32], shape: 3): the target position the camera suppose to look at
|
||||
target_to_cam_dir (NDArray[float32], shape: 3): offset direction from target to camera
|
||||
look_distance (float, optional): length of offset vector from target to camera.
|
||||
|
||||
Returns:
|
||||
NDArray[float32]: shape: (4, 4), pose (cam2world) matrix
|
||||
"""
|
||||
|
||||
norm=np.linalg.norm(target_to_cam_offset)
|
||||
if norm==0:
|
||||
norm=np.finfo(np.float32).eps
|
||||
target_to_cam_offset = look_distance * target_to_cam_offset / norm
|
||||
|
||||
campos = target_to_cam_offset + target # [3]
|
||||
T = np.eye(4, dtype=np.float32)
|
||||
T[:3, :3] = look_at(campos, target, opengl)
|
||||
T[:3, 3] = campos
|
||||
return T
|
||||
|
||||
class OrbitCamera:
|
||||
def __init__(self, W, H, r=2, fovy=60, near=0.01, far=100):
|
||||
self.W = W
|
||||
self.H = H
|
||||
self.radius = r # camera distance from center
|
||||
self.fovy = np.deg2rad(fovy) # deg 2 rad
|
||||
self.near = near
|
||||
self.far = far
|
||||
self.center = np.array([0, 0, 0], dtype=np.float32) # look at this point
|
||||
self.rot = R.from_matrix(np.eye(3))
|
||||
self.up = np.array([0, 1, 0], dtype=np.float32) # need to be normalized!
|
||||
|
||||
@property
|
||||
def fovx(self):
|
||||
return 2 * np.arctan(np.tan(self.fovy / 2) * self.W / self.H)
|
||||
|
||||
@property
|
||||
def campos(self):
|
||||
return self.pose[:3, 3]
|
||||
|
||||
# pose (c2w)
|
||||
@property
|
||||
def pose(self):
|
||||
# first move camera to radius
|
||||
res = np.eye(4, dtype=np.float32)
|
||||
res[2, 3] = self.radius # opengl convention...
|
||||
# rotate
|
||||
rot = np.eye(4, dtype=np.float32)
|
||||
rot[:3, :3] = self.rot.as_matrix()
|
||||
res = rot @ res
|
||||
# translate
|
||||
res[:3, 3] -= self.center
|
||||
return res
|
||||
|
||||
# view (w2c)
|
||||
@property
|
||||
def view(self):
|
||||
return np.linalg.inv(self.pose)
|
||||
|
||||
# projection (perspective)
|
||||
@property
|
||||
def perspective(self):
|
||||
y = np.tan(self.fovy / 2)
|
||||
aspect = self.W / self.H
|
||||
return np.array(
|
||||
[
|
||||
[1 / (y * aspect), 0, 0, 0],
|
||||
[0, -1 / y, 0, 0],
|
||||
[
|
||||
0,
|
||||
0,
|
||||
-(self.far + self.near) / (self.far - self.near),
|
||||
-(2 * self.far * self.near) / (self.far - self.near),
|
||||
],
|
||||
[0, 0, -1, 0],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
# intrinsics
|
||||
@property
|
||||
def intrinsics(self):
|
||||
focal = self.H / (2 * np.tan(self.fovy / 2))
|
||||
return np.array([focal, focal, self.W // 2, self.H // 2], dtype=np.float32)
|
||||
|
||||
@property
|
||||
def mvp(self):
|
||||
return self.perspective @ np.linalg.inv(self.pose) # [4, 4]
|
||||
|
||||
def orbit(self, dx, dy):
|
||||
# rotate along camera up/side axis!
|
||||
side = self.rot.as_matrix()[:3, 0]
|
||||
rotvec_x = self.up * np.radians(-0.05 * dx)
|
||||
rotvec_y = side * np.radians(-0.05 * dy)
|
||||
self.rot = R.from_rotvec(rotvec_x) * R.from_rotvec(rotvec_y) * self.rot
|
||||
|
||||
def scale(self, delta):
|
||||
self.radius *= 1.1 ** (-delta)
|
||||
|
||||
def pan(self, dx, dy, dz=0):
|
||||
# pan in camera coordinate system (careful on the sensitivity!)
|
||||
self.center += 0.0005 * self.rot.as_matrix()[:3, :3] @ np.array([-dx, -dy, dz])
|
||||
|
||||
def calculate_fovX(H, W, fovy):
|
||||
return 2 * np.arctan(np.tan(fovy / 2) * W / H)
|
||||
|
||||
def get_projection_matrix(znear, zfar, fovX, fovY, z_sign=1.0):
|
||||
tanHalfFovY = math.tan((fovY / 2))
|
||||
tanHalfFovX = math.tan((fovX / 2))
|
||||
|
||||
P = torch.zeros(4, 4)
|
||||
|
||||
P[0, 0] = 1 / tanHalfFovX
|
||||
P[1, 1] = 1 / tanHalfFovY
|
||||
P[3, 2] = z_sign
|
||||
P[2, 2] = z_sign * zfar / (zfar - znear)
|
||||
P[2, 3] = -(zfar * znear) / (zfar - znear)
|
||||
return P
|
||||
|
||||
|
||||
class MiniCam:
|
||||
def __init__(self, c2w, width, height, fovy, fovx, znear, zfar, projection_matrix=None):
|
||||
# c2w (pose) should be in NeRF convention.
|
||||
|
||||
self.image_width = width
|
||||
self.image_height = height
|
||||
self.FoVy = fovy
|
||||
self.FoVx = fovx
|
||||
self.znear = znear
|
||||
self.zfar = zfar
|
||||
|
||||
w2c = np.linalg.inv(c2w)
|
||||
|
||||
# rectify...
|
||||
w2c[1:3, :3] *= -1
|
||||
w2c[:3, 3] *= -1
|
||||
|
||||
self.world_view_transform = torch.tensor(w2c).transpose(0, 1).cuda()
|
||||
self.projection_matrix = (
|
||||
get_projection_matrix(
|
||||
znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy
|
||||
)
|
||||
.transpose(0, 1)
|
||||
.cuda()
|
||||
) if projection_matrix is None else projection_matrix
|
||||
self.full_proj_transform = self.world_view_transform @ self.projection_matrix
|
||||
self.camera_center = -torch.tensor(c2w[:3, 3]).cuda()
|
||||
|
||||
class BaseCameraController(ABC):
|
||||
def __init__(self, renderer, cam_size_W, cam_size_H, reference_orbit_camera_fovy, invert_bg_prob=1.0, static_bg=None, device='cuda'):
|
||||
self.device = torch.device(device)
|
||||
|
||||
self.renderer = renderer
|
||||
self.cam = OrbitCamera(cam_size_W, cam_size_H, fovy=reference_orbit_camera_fovy)
|
||||
|
||||
self.invert_bg_prob = invert_bg_prob
|
||||
self.black_bg = torch.tensor([0, 0, 0], dtype=torch.float32, device=self.device)
|
||||
self.white_bg = torch.tensor([1, 1, 1], dtype=torch.float32, device=self.device)
|
||||
self.static_bg = None if static_bg is None else torch.tensor(static_bg, dtype=torch.float32, device=self.device)
|
||||
|
||||
self.post_init()
|
||||
|
||||
super().__init__()
|
||||
|
||||
def post_init(self):
|
||||
# Calls after default initialize at the end of __init__()
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_render_result(self, render_pose, bg_color, **kwargs):
|
||||
pass
|
||||
|
||||
def render_at_pose(self, cam_pose, **kwargs):
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = cam_pose
|
||||
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
render_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
|
||||
if self.static_bg is None:
|
||||
bg_color = self.white_bg if np.random.rand() > self.invert_bg_prob else self.black_bg
|
||||
else:
|
||||
bg_color = self.static_bg
|
||||
|
||||
return self.get_render_result(render_pose, bg_color, **kwargs)
|
||||
|
||||
def render_all_pose(self, all_cam_poses, **kwargs):
|
||||
all_rendered_images, all_rendered_masks = [], []
|
||||
extra_outputs = {}
|
||||
for cam_pose in all_cam_poses:
|
||||
out = self.render_at_pose(cam_pose, **kwargs)
|
||||
|
||||
image = out["image"] # [3, H, W] in [0, 1]
|
||||
mask = out["alpha"] # [1, H, W] in [0, 1]
|
||||
|
||||
all_rendered_images.append(image)
|
||||
all_rendered_masks.append(mask)
|
||||
|
||||
for k in out:
|
||||
if k not in extra_outputs:
|
||||
extra_outputs[k] = []
|
||||
extra_outputs[k].append(out[k])
|
||||
|
||||
for k in extra_outputs:
|
||||
extra_outputs[k] = torch.stack(extra_outputs[k], dim=0)
|
||||
|
||||
# [Number of Poses, 3, H, W], [Number of Poses, 1, H, W] both in [0, 1]
|
||||
return torch.stack(all_rendered_images, dim=0), torch.stack(all_rendered_masks, dim=0), extra_outputs
|
||||
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
from os import listdir
|
||||
from os.path import isfile, join, exists, dirname
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from shared_utils.log_utils import cstr
|
||||
|
||||
def get_parent_dirpath_n_level_up(abs_path, n=1):
|
||||
for i in range(n):
|
||||
abs_path = dirname(abs_path)
|
||||
return abs_path
|
||||
|
||||
def get_persistent_directory(folder_name):
|
||||
if sys.platform == "win32":
|
||||
folder = join(os.path.expanduser("~"), "AppData", "Local", folder_name)
|
||||
else:
|
||||
folder = join(os.path.expanduser("~"), "." + folder_name)
|
||||
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
return folder
|
||||
|
||||
def parse_save_filename(save_path, output_directory, supported_extensions, class_name):
|
||||
|
||||
folder_path, filename = os.path.split(save_path)
|
||||
filename, file_extension = os.path.splitext(filename)
|
||||
if file_extension.lower() in supported_extensions:
|
||||
if not os.path.isabs(save_path):
|
||||
folder_path = join(output_directory, folder_path)
|
||||
|
||||
os.makedirs(folder_path, exist_ok=True)
|
||||
|
||||
# replace time date format to current time
|
||||
now = datetime.now() # current date and time
|
||||
all_date_format = ["%Y", "%m", "%d", "%H", "%M", "%S", "%f"]
|
||||
for date_format in all_date_format:
|
||||
if date_format in filename:
|
||||
filename = filename.replace(date_format, now.strftime(date_format))
|
||||
|
||||
save_path = join(folder_path, filename) + file_extension
|
||||
cstr(f"[{class_name}] Saving model to {save_path}").msg.print()
|
||||
return save_path
|
||||
else:
|
||||
cstr(f"[{class_name}] File name {filename} does not end with supported file extensions: {supported_extensions}").error.print()
|
||||
|
||||
return None
|
||||
|
||||
def get_list_filenames(directory, extension_filter=None, recursive=False):
|
||||
"""
|
||||
Recursively finds files with specified extensions in a directory and returns relative paths.
|
||||
|
||||
Args:
|
||||
directory (str): The directory path to search.
|
||||
extension_filter (list): List of file extensions (e.g., ['.txt', '.csv']).
|
||||
|
||||
Returns:
|
||||
list: List of relative file paths matching the specified extensions.
|
||||
"""
|
||||
if exists(directory):
|
||||
if recursive:
|
||||
result = []
|
||||
for root, _, files in os.walk(directory):
|
||||
for item in files:
|
||||
if extension_filter is None or os.path.splitext(item)[1].lower() in extension_filter:
|
||||
relative_path = os.path.relpath(os.path.join(root, item), directory)
|
||||
result.append(relative_path)
|
||||
return result
|
||||
else:
|
||||
return [f for f in listdir(directory) if isfile(join(directory, f)) and (extension_filter is None or f.lower().endswith(extension_filter))]
|
||||
else:
|
||||
return []
|
||||
|
||||
# Download pre-trained model if it not exist locally
|
||||
def resume_or_download_model_from_hf(checkpoints_dir_abs, repo_id, model_name, class_name="", repo_type="model"):
|
||||
|
||||
ckpt_path = os.path.join(checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{class_name}] can't find checkpoint {ckpt_path}, will download it from repo {repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=repo_id, local_dir=checkpoints_dir_abs, filename=model_name, repo_type=repo_type)
|
||||
|
||||
return ckpt_path
|
||||
@@ -0,0 +1,198 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms.functional as TF
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
def prepare_torch_img(img, size_H, size_W, device="cuda", keep_shape=False):
|
||||
# [N, H, W, C] -> [N, C, H, W]
|
||||
img_new = img.permute(0, 3, 1, 2).to(device)
|
||||
img_new = F.interpolate(img_new, (size_H, size_W), mode="bilinear", align_corners=False).contiguous()
|
||||
if keep_shape:
|
||||
img_new = img_new.permute(0, 2, 3, 1)
|
||||
return img_new
|
||||
|
||||
def torch_imgs_to_pils(images, masks=None, alpha_min=0.1):
|
||||
"""
|
||||
images (torch): [N, H, W, C] or [H, W, C]
|
||||
masks (torch): [N, H, W] or [H, W]
|
||||
"""
|
||||
if len(images.shape) == 3:
|
||||
images = images.unsqueeze(0)
|
||||
|
||||
if masks is not None:
|
||||
if len(masks.shape) == 2:
|
||||
masks = masks.unsqueeze(0)
|
||||
|
||||
inv_mask_index = masks < alpha_min
|
||||
images[inv_mask_index] = 0.
|
||||
|
||||
masks = masks.unsqueeze(3)
|
||||
images = torch.cat((images, masks), dim=3)
|
||||
mode="RGBA"
|
||||
else:
|
||||
mode="RGB"
|
||||
|
||||
pil_image_list = [Image.fromarray((images[i].detach().cpu().numpy() * 255).astype(np.uint8), mode=mode) for i in range(images.shape[0])]
|
||||
|
||||
return pil_image_list
|
||||
|
||||
def troch_image_dilate(img):
|
||||
"""
|
||||
Remove thin seams on generated texture
|
||||
img (torch): [H, W, C]
|
||||
"""
|
||||
import cv2
|
||||
img = np.asarray(img.cpu().numpy(), dtype=np.float32)
|
||||
img = img * 255
|
||||
img = img.clip(0, 255)
|
||||
mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True)
|
||||
mask = (mask <= 3.0).astype(np.float32)
|
||||
kernel = np.ones((3, 3), 'uint8')
|
||||
dilate_img = cv2.dilate(img, kernel, iterations=1)
|
||||
img = img * (1 - mask) + dilate_img * mask
|
||||
img = (img.clip(0, 255) / 255).astype(np.float32)
|
||||
return torch.from_numpy(img)
|
||||
|
||||
def pils_to_torch_imgs(pils: Union[Image.Image, List[Image.Image]], device="cuda", force_rgb=True):
|
||||
if isinstance(pils, Image.Image):
|
||||
pils = [pils]
|
||||
|
||||
images = []
|
||||
for pil in pils:
|
||||
if pil.mode == "RGBA" and force_rgb:
|
||||
pil = pil.convert('RGB')
|
||||
|
||||
images.append(TF.to_tensor(pil).permute(1, 2, 0))
|
||||
|
||||
images = torch.stack(images, dim=0).to(device)
|
||||
|
||||
return images
|
||||
|
||||
def pils_rgba_to_rgb(pils: Union[Image.Image, List[Image.Image]], bkgd="WHITE"):
|
||||
if isinstance(pils, Image.Image):
|
||||
pils = [pils]
|
||||
|
||||
rgbs = []
|
||||
for pil in pils:
|
||||
if pil.mode == 'RGBA':
|
||||
new_image = Image.new("RGBA", pil.size, bkgd)
|
||||
new_image.paste(pil, (0, 0), pil)
|
||||
rgbs.append(new_image.convert('RGB'))
|
||||
else:
|
||||
rgbs.append(pil)
|
||||
|
||||
return rgbs
|
||||
|
||||
def pil_split_image(image, rows=None, cols=None):
|
||||
"""
|
||||
inverse function of make_image_grid
|
||||
"""
|
||||
# image is in square
|
||||
if rows is None and cols is None:
|
||||
# image.size [W, H]
|
||||
rows = 1
|
||||
cols = image.size[0] // image.size[1]
|
||||
assert cols * image.size[1] == image.size[0]
|
||||
subimg_size = image.size[1]
|
||||
elif rows is None:
|
||||
subimg_size = image.size[0] // cols
|
||||
rows = image.size[1] // subimg_size
|
||||
assert rows * subimg_size == image.size[1]
|
||||
elif cols is None:
|
||||
subimg_size = image.size[1] // rows
|
||||
cols = image.size[0] // subimg_size
|
||||
assert cols * subimg_size == image.size[0]
|
||||
else:
|
||||
subimg_size = image.size[1] // rows
|
||||
assert cols * subimg_size == image.size[0]
|
||||
subimgs = []
|
||||
for i in range(rows):
|
||||
for j in range(cols):
|
||||
subimg = image.crop((j*subimg_size, i*subimg_size, (j+1)*subimg_size, (i+1)*subimg_size))
|
||||
subimgs.append(subimg)
|
||||
return subimgs
|
||||
|
||||
def pil_make_image_grid(images, rows=None, cols=None):
|
||||
if rows is None and cols is None:
|
||||
rows = 1
|
||||
cols = len(images)
|
||||
if rows is None:
|
||||
rows = len(images) // cols
|
||||
if len(images) % cols != 0:
|
||||
rows += 1
|
||||
if cols is None:
|
||||
cols = len(images) // rows
|
||||
if len(images) % rows != 0:
|
||||
cols += 1
|
||||
total_imgs = rows * cols
|
||||
if total_imgs > len(images):
|
||||
images += [Image.new(images[0].mode, images[0].size) for _ in range(total_imgs - len(images))]
|
||||
|
||||
w, h = images[0].size
|
||||
grid = Image.new(images[0].mode, size=(cols * w, rows * h))
|
||||
|
||||
for i, img in enumerate(images):
|
||||
grid.paste(img, box=(i % cols * w, i // cols * h))
|
||||
return grid
|
||||
|
||||
def pils_erode_masks(mask_list):
|
||||
out_mask_list = []
|
||||
for idx, mask in enumerate(mask_list):
|
||||
arr = np.array(mask)
|
||||
alpha = (arr[:, :, 3] > 127).astype(np.uint8)
|
||||
# erode 1px
|
||||
import cv2
|
||||
alpha = cv2.erode(alpha, np.ones((3, 3), np.uint8), iterations=1)
|
||||
alpha = (alpha * 255).astype(np.uint8)
|
||||
out_mask_list.append(Image.fromarray(alpha[:, :, None]))
|
||||
|
||||
return out_mask_list
|
||||
|
||||
def pils_resize_foreground(
|
||||
pils: Union[Image.Image, List[Image.Image]],
|
||||
ratio: float,
|
||||
) -> List[Image.Image]:
|
||||
if isinstance(pils, Image.Image):
|
||||
pils = [pils]
|
||||
|
||||
new_pils = []
|
||||
for image in pils:
|
||||
image = np.array(image)
|
||||
assert image.shape[-1] == 4
|
||||
alpha = np.where(image[..., 3] > 0)
|
||||
y1, y2, x1, x2 = (
|
||||
alpha[0].min(),
|
||||
alpha[0].max(),
|
||||
alpha[1].min(),
|
||||
alpha[1].max(),
|
||||
)
|
||||
# crop the foreground
|
||||
fg = image[y1:y2, x1:x2]
|
||||
# pad to square
|
||||
size = max(fg.shape[0], fg.shape[1])
|
||||
ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
|
||||
ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
|
||||
new_image = np.pad(
|
||||
fg,
|
||||
((ph0, ph1), (pw0, pw1), (0, 0)),
|
||||
mode="constant",
|
||||
constant_values=((0, 0), (0, 0), (0, 0)),
|
||||
)
|
||||
|
||||
# compute padding according to the ratio
|
||||
new_size = int(new_image.shape[0] / ratio)
|
||||
# pad to size, double side
|
||||
ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
|
||||
ph1, pw1 = new_size - size - ph0, new_size - size - pw0
|
||||
new_image = np.pad(
|
||||
new_image,
|
||||
((ph0, ph1), (pw0, pw1), (0, 0)),
|
||||
mode="constant",
|
||||
constant_values=((0, 0), (0, 0), (0, 0)),
|
||||
)
|
||||
new_image = Image.fromarray(new_image, mode="RGBA")
|
||||
new_pils.append(new_image)
|
||||
|
||||
return new_pils
|
||||
@@ -0,0 +1,105 @@
|
||||
import logging
|
||||
import sys
|
||||
|
||||
class cstr(str):
|
||||
# Modified from: WAS Node Suite
|
||||
class color:
|
||||
END = '\33[0m'
|
||||
BOLD = '\33[1m'
|
||||
ITALIC = '\33[3m'
|
||||
UNDERLINE = '\33[4m'
|
||||
BLINK = '\33[5m'
|
||||
BLINK2 = '\33[6m'
|
||||
SELECTED = '\33[7m'
|
||||
|
||||
BLACK = '\33[30m'
|
||||
RED = '\33[31m'
|
||||
GREEN = '\33[32m'
|
||||
YELLOW = '\33[33m'
|
||||
BLUE = '\33[34m'
|
||||
VIOLET = '\33[35m'
|
||||
BEIGE = '\33[36m'
|
||||
WHITE = '\33[37m'
|
||||
|
||||
BLACKBG = '\33[40m'
|
||||
REDBG = '\33[41m'
|
||||
GREENBG = '\33[42m'
|
||||
YELLOWBG = '\33[43m'
|
||||
BLUEBG = '\33[44m'
|
||||
VIOLETBG = '\33[45m'
|
||||
BEIGEBG = '\33[46m'
|
||||
WHITEBG = '\33[47m'
|
||||
|
||||
GREY = '\33[90m'
|
||||
LIGHTRED = '\33[91m'
|
||||
LIGHTGREEN = '\33[92m'
|
||||
LIGHTYELLOW = '\33[93m'
|
||||
LIGHTBLUE = '\33[94m'
|
||||
LIGHTVIOLET = '\33[95m'
|
||||
LIGHTBEIGE = '\33[96m'
|
||||
LIGHTWHITE = '\33[97m'
|
||||
|
||||
GREYBG = '\33[100m'
|
||||
LIGHTREDBG = '\33[101m'
|
||||
LIGHTGREENBG = '\33[102m'
|
||||
LIGHTYELLOWBG = '\33[103m'
|
||||
LIGHTBLUEBG = '\33[104m'
|
||||
LIGHTVIOLETBG = '\33[105m'
|
||||
LIGHTBEIGEBG = '\33[106m'
|
||||
LIGHTWHITEBG = '\33[107m'
|
||||
|
||||
@staticmethod
|
||||
def add_code(name, code):
|
||||
if not hasattr(cstr.color, name.upper()):
|
||||
setattr(cstr.color, name.upper(), code)
|
||||
else:
|
||||
raise ValueError(f"'cstr' object already contains a code with the name '{name}'.")
|
||||
|
||||
def __new__(cls, text):
|
||||
return super().__new__(cls, text)
|
||||
|
||||
def __getattr__(self, attr):
|
||||
if attr.lower().startswith("_cstr"):
|
||||
code = getattr(self.color, attr.upper().lstrip("_cstr"))
|
||||
modified_text = self.replace(f"__{attr[1:]}__", f"{code}")
|
||||
return cstr(modified_text)
|
||||
elif attr.upper() in dir(self.color):
|
||||
code = getattr(self.color, attr.upper())
|
||||
modified_text = f"{code}{self}{self.color.END}"
|
||||
return cstr(modified_text)
|
||||
elif attr.lower() in dir(cstr):
|
||||
return getattr(cstr, attr.lower())
|
||||
else:
|
||||
raise AttributeError(f"'cstr' object has no attribute '{attr}'")
|
||||
|
||||
def print(self, **kwargs):
|
||||
print(self, **kwargs)
|
||||
|
||||
#! MESSAGE TEMPLATES
|
||||
cstr.color.add_code("msg", f"{cstr.color.BLUE}[Comfy3D] {cstr.color.END}")
|
||||
cstr.color.add_code("warning", f"{cstr.color.LIGHTYELLOW}[Comfy3D] [WARNING] {cstr.color.END}")
|
||||
cstr.color.add_code("error", f"{cstr.color.RED}[Comfy3D] [ERROR] {cstr.color.END}")
|
||||
|
||||
class WarningFilter(logging.Filter):
|
||||
def filter(self, record):
|
||||
if record.levelno == logging.WARNING:
|
||||
record.msg = f"Warn!: {record.msg}"
|
||||
return True
|
||||
|
||||
def create_handler(stream, levels, formatter):
|
||||
handler = logging.StreamHandler(stream)
|
||||
handler.setLevel(min(levels))
|
||||
handler.addFilter(lambda record: record.levelno in levels)
|
||||
handler.addFilter(WarningFilter()) # Apply the custom filter
|
||||
handler.setFormatter(formatter)
|
||||
|
||||
return handler
|
||||
def setup_logger(logger_name, level, stdout_levels, stderr_levels, formatter):
|
||||
logger = logging.getLogger(logger_name)
|
||||
for handler in logger.handlers[:]:
|
||||
logger.removeHandler(handler)
|
||||
logger.setLevel(level)
|
||||
stdout_handler = create_handler(sys.stdout, stdout_levels, formatter)
|
||||
stderr_handler = create_handler(sys.stderr, stderr_levels, formatter)
|
||||
logger.addHandler(stdout_handler)
|
||||
logger.addHandler(stderr_handler)
|
||||
@@ -0,0 +1,118 @@
|
||||
# Copyright 2021 The PlenOctree Authors.
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
|
||||
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
||||
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
||||
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
||||
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
||||
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
# POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
import torch
|
||||
|
||||
C0 = 0.28209479177387814
|
||||
C1 = 0.4886025119029199
|
||||
C2 = [
|
||||
1.0925484305920792,
|
||||
-1.0925484305920792,
|
||||
0.31539156525252005,
|
||||
-1.0925484305920792,
|
||||
0.5462742152960396
|
||||
]
|
||||
C3 = [
|
||||
-0.5900435899266435,
|
||||
2.890611442640554,
|
||||
-0.4570457994644658,
|
||||
0.3731763325901154,
|
||||
-0.4570457994644658,
|
||||
1.445305721320277,
|
||||
-0.5900435899266435
|
||||
]
|
||||
C4 = [
|
||||
2.5033429417967046,
|
||||
-1.7701307697799304,
|
||||
0.9461746957575601,
|
||||
-0.6690465435572892,
|
||||
0.10578554691520431,
|
||||
-0.6690465435572892,
|
||||
0.47308734787878004,
|
||||
-1.7701307697799304,
|
||||
0.6258357354491761,
|
||||
]
|
||||
|
||||
|
||||
def eval_sh(deg, sh, dirs):
|
||||
"""
|
||||
Evaluate spherical harmonics at unit directions
|
||||
using hardcoded SH polynomials.
|
||||
Works with torch/np/jnp.
|
||||
... Can be 0 or more batch dimensions.
|
||||
Args:
|
||||
deg: int SH deg. Currently, 0-3 supported
|
||||
sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
|
||||
dirs: jnp.ndarray unit directions [..., 3]
|
||||
Returns:
|
||||
[..., C]
|
||||
"""
|
||||
assert deg <= 4 and deg >= 0
|
||||
coeff = (deg + 1) ** 2
|
||||
assert sh.shape[-1] >= coeff
|
||||
|
||||
result = C0 * sh[..., 0]
|
||||
if deg > 0:
|
||||
x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
|
||||
result = (result -
|
||||
C1 * y * sh[..., 1] +
|
||||
C1 * z * sh[..., 2] -
|
||||
C1 * x * sh[..., 3])
|
||||
|
||||
if deg > 1:
|
||||
xx, yy, zz = x * x, y * y, z * z
|
||||
xy, yz, xz = x * y, y * z, x * z
|
||||
result = (result +
|
||||
C2[0] * xy * sh[..., 4] +
|
||||
C2[1] * yz * sh[..., 5] +
|
||||
C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
|
||||
C2[3] * xz * sh[..., 7] +
|
||||
C2[4] * (xx - yy) * sh[..., 8])
|
||||
|
||||
if deg > 2:
|
||||
result = (result +
|
||||
C3[0] * y * (3 * xx - yy) * sh[..., 9] +
|
||||
C3[1] * xy * z * sh[..., 10] +
|
||||
C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
|
||||
C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
|
||||
C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
|
||||
C3[5] * z * (xx - yy) * sh[..., 14] +
|
||||
C3[6] * x * (xx - 3 * yy) * sh[..., 15])
|
||||
|
||||
if deg > 3:
|
||||
result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
|
||||
C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
|
||||
C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
|
||||
C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
|
||||
C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
|
||||
C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
|
||||
C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
|
||||
C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
|
||||
C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
|
||||
return result
|
||||
|
||||
def RGB2SH(rgb):
|
||||
return (rgb - 0.5) / C0
|
||||
|
||||
def SH2RGB(sh):
|
||||
return sh * C0 + 0.5
|
||||
Reference in New Issue
Block a user