222 lines
7.0 KiB
Python
222 lines
7.0 KiB
Python
import torch
|
|
import numpy as np
|
|
from scipy.fft import dct, idct
|
|
|
|
# ==============================================================================
|
|
# HELPER FUNCTIONS & ALGORITHMS
|
|
# ==============================================================================
|
|
|
|
def get_gradients_from_normal(normal_map):
|
|
"""Convert RGB normal map to gradients (dz/dx, dz/dy)."""
|
|
norm = normal_map * 2.0 - 1.0
|
|
nx = norm[:, :, 0]
|
|
ny = norm[:, :, 1]
|
|
nz = norm[:, :, 2]
|
|
|
|
# Handle tiny nz values
|
|
eps = 1e-6
|
|
nz = np.where(np.abs(nz) < eps, eps * np.sign(nz), nz)
|
|
|
|
p = -nx / nz
|
|
q = -ny / nz
|
|
return p, q
|
|
|
|
def normalize_output(depth_map):
|
|
"""Normalize depth map to [0, 1] range."""
|
|
d_min = np.min(depth_map)
|
|
d_max = np.max(depth_map)
|
|
if d_max - d_min < 1e-6:
|
|
return np.zeros_like(depth_map)
|
|
return (depth_map - d_min) / (d_max - d_min)
|
|
|
|
def frankot_chellappa(p, q):
|
|
"""Project gradients into Fourier domain, integrate, and inverse transform."""
|
|
rows, cols = p.shape
|
|
fp = np.fft.fft2(p)
|
|
fq = np.fft.fft2(q)
|
|
|
|
u, v = np.meshgrid(np.fft.fftfreq(cols), np.fft.fftfreq(rows))
|
|
|
|
numerator = -1j * (u * fp + v * fq)
|
|
denominator = (u**2 + v**2) + 1e-12
|
|
|
|
fz = numerator / denominator
|
|
fz[0, 0] = 0
|
|
|
|
z = np.fft.ifft2(fz).real
|
|
return z
|
|
|
|
def poisson_reconstruction(p, q):
|
|
"""Solves Poisson equation: Laplacian(Z) = div(Gradients) using DCT."""
|
|
rows, cols = p.shape
|
|
dx_p = np.gradient(p, axis=1)
|
|
dy_q = np.gradient(q, axis=0)
|
|
divergence = dx_p + dy_q
|
|
|
|
dct_div = dct(dct(divergence, axis=0, norm='ortho'), axis=1, norm='ortho')
|
|
|
|
x_coords = np.arange(cols)
|
|
y_coords = np.arange(rows)
|
|
xv, yv = np.meshgrid(x_coords, y_coords)
|
|
|
|
eigenvalues = (2 * np.cos(np.pi * xv / cols) - 2) + \
|
|
(2 * np.cos(np.pi * yv / rows) - 2)
|
|
|
|
eigenvalues[0, 0] = 1.0
|
|
dct_z = dct_div / eigenvalues
|
|
dct_z[0, 0] = 0.0
|
|
|
|
z = idct(idct(dct_z, axis=1, norm='ortho'), axis=0, norm='ortho')
|
|
return z
|
|
|
|
# ==============================================================================
|
|
# COMFYUI NODES
|
|
# ==============================================================================
|
|
|
|
class NormalToDepthFrankotChellappa:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"normal_map": ("IMAGE",),
|
|
"invert_output": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("depth_map",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "DepthSolvers"
|
|
|
|
def process(self, normal_map, invert_output):
|
|
nm_np = normal_map.cpu().numpy()
|
|
batch_size, h, w, c = nm_np.shape
|
|
output_batch = []
|
|
|
|
for i in range(batch_size):
|
|
p, q = get_gradients_from_normal(nm_np[i])
|
|
z = frankot_chellappa(p, q)
|
|
z_norm = normalize_output(z)
|
|
if invert_output:
|
|
z_norm = 1.0 - z_norm
|
|
z_out = np.stack([z_norm, z_norm, z_norm], axis=-1)
|
|
output_batch.append(z_out)
|
|
|
|
return (torch.from_numpy(np.array(output_batch)).float(),)
|
|
|
|
|
|
class NormalToDepthPoisson:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"normal_map": ("IMAGE",),
|
|
"invert_output": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("depth_map",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "DepthSolvers"
|
|
|
|
def process(self, normal_map, invert_output):
|
|
nm_np = normal_map.cpu().numpy()
|
|
batch_size, h, w, c = nm_np.shape
|
|
output_batch = []
|
|
|
|
for i in range(batch_size):
|
|
p, q = get_gradients_from_normal(nm_np[i])
|
|
z = poisson_reconstruction(p, q)
|
|
z_norm = normalize_output(z)
|
|
if invert_output:
|
|
z_norm = 1.0 - z_norm
|
|
z_out = np.stack([z_norm, z_norm, z_norm], axis=-1)
|
|
output_batch.append(z_out)
|
|
|
|
return (torch.from_numpy(np.array(output_batch)).float(),)
|
|
|
|
|
|
class DepthMapMathCombiner:
|
|
"""
|
|
Combines two depth maps using a mathematical expression.
|
|
Inputs: X (depth_map_1), Y (depth_map_2), a, b
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"depth_map_x": ("IMAGE",),
|
|
"depth_map_y": ("IMAGE",),
|
|
"expression": ("STRING", {"default": "a * X + b * Y", "multiline": False}),
|
|
"a": ("FLOAT", {"default": 1.0, "step": 0.01, "min": -100.0, "max": 100.0}),
|
|
"b": ("FLOAT", {"default": 1.0, "step": 0.01, "min": -100.0, "max": 100.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("combined_depth",)
|
|
FUNCTION = "combine"
|
|
CATEGORY = "DepthSolvers"
|
|
|
|
def combine(self, depth_map_x, depth_map_y, expression, a, b):
|
|
# Ensure dimensions match for broadcasting if possible, or raise error
|
|
if depth_map_x.shape != depth_map_y.shape:
|
|
# Basic safety: if shapes don't match, try to use X's shape
|
|
# In a real scenario, you might want to resize Y to X here.
|
|
# For now, we assume user provides matching dimensions.
|
|
pass
|
|
|
|
# Define the context for the evaluation
|
|
# We use Torch math functions for speed
|
|
math_context = {
|
|
"X": depth_map_x,
|
|
"Y": depth_map_y,
|
|
"a": a,
|
|
"b": b,
|
|
"torch": torch,
|
|
"log": torch.log,
|
|
"exp": torch.exp,
|
|
"sqrt": torch.sqrt,
|
|
"sin": torch.sin,
|
|
"cos": torch.cos,
|
|
"abs": torch.abs,
|
|
"max": torch.max,
|
|
"min": torch.min,
|
|
"pow": torch.pow
|
|
}
|
|
|
|
try:
|
|
# Evaluate the expression
|
|
# Example: "a * X + b * torch.log(Y + 0.001)"
|
|
result = eval(expression, {"__builtins__": {}}, math_context)
|
|
|
|
# Handle result if it's a scalar (rare, but possible)
|
|
if isinstance(result, (int, float)):
|
|
result = torch.full_like(depth_map_x, result)
|
|
|
|
# Clamp result to 0-1 to stay valid as an image (Optional, but recommended for depth)
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
return (result,)
|
|
|
|
except Exception as e:
|
|
print(f"[DepthMapMathCombiner] Error evaluating expression: {e}")
|
|
# Return X as fallback in case of error
|
|
return (depth_map_x,)
|
|
|
|
# ==============================================================================
|
|
# MAPPINGS
|
|
# ==============================================================================
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FrankotChellappa": NormalToDepthFrankotChellappa,
|
|
"PoissonReconstruction": NormalToDepthPoisson,
|
|
"DepthMathCombiner": DepthMapMathCombiner
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FrankotChellappa": "Normal to Depth (Frankot-Chellappa)",
|
|
"PoissonReconstruction": "Normal to Depth (Poisson)",
|
|
"DepthMathCombiner": "Depth Map Math Combiner"
|
|
} |