254 lines
18 KiB
Python
254 lines
18 KiB
Python
import torch
|
|
import numpy as np
|
|
import cv2
|
|
import folder_paths
|
|
import os
|
|
from PIL import Image
|
|
|
|
# Imports pour le serveur sécurisé
|
|
import http.server
|
|
import socketserver
|
|
import threading
|
|
|
|
# Le bloc manquant, remis à sa place
|
|
try:
|
|
from scipy.ndimage import distance_transform_edt
|
|
SCIPY_AVAILABLE = True
|
|
except ImportError:
|
|
SCIPY_AVAILABLE = False
|
|
|
|
# Variable globale pour s'assurer que le serveur n'est démarré qu'une seule fois
|
|
preview_server_thread = None
|
|
PREVIEW_SERVER_PORT = 8189 # Port pour notre serveur
|
|
|
|
def start_preview_server():
|
|
global preview_server_thread
|
|
if preview_server_thread is None or not preview_server_thread.is_alive():
|
|
class SecureHandler(http.server.SimpleHTTPRequestHandler):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, directory=folder_paths.get_temp_directory(), **kwargs)
|
|
def do_GET(self):
|
|
if self.path == '/' or self.path.endswith('/'):
|
|
self.send_error(403, "Directory listing is not allowed")
|
|
return
|
|
super().do_GET()
|
|
def log_message(self, format, *args):
|
|
return
|
|
address = ("127.0.0.1", PREVIEW_SERVER_PORT)
|
|
httpd = socketserver.TCPServer(address, SecureHandler)
|
|
thread = threading.Thread(target=httpd.serve_forever)
|
|
thread.daemon = True
|
|
thread.start()
|
|
preview_server_thread = thread
|
|
print(f"[Relight Node] Serveur d'aperçu sécurisé démarré sur http://127.0.0.1:{PREVIEW_SERVER_PORT}")
|
|
|
|
|
|
class RelightNode:
|
|
OUTPUT_NODE = True
|
|
|
|
LIGHT_TYPES = ["Point", "Neon"]
|
|
DEBUG_VIEWS = ["Final Image", "Lighting Map Only", "Mask Lighting Only", "Mask Virtual Normals", "Light Shape Viz"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
if not SCIPY_AVAILABLE:
|
|
return {"required": {"error": ("STRING", {"default": "ERREUR: SciPy n'est pas installé. Exécutez 'pip install scipy' et redémarrez ComfyUI.", "multiline": True})}}
|
|
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"normal_map": ("IMAGE",),
|
|
"depth_map": ("IMAGE",),
|
|
"debug_view": (s.DEBUG_VIEWS, ),
|
|
"preserve_color": ("BOOLEAN", {"default": True}),
|
|
"ambient_light": ("FLOAT", {"default": -0.20, "min": -1.0, "max": 1.0, "step": 0.05}),
|
|
"gamma": ("FLOAT", {"default": 5.5, "min": 0.1, "max": 20.0, "step": 0.1}),
|
|
"light_color": ("STRING", {"default": "#FFFFFF"}),
|
|
"depth_scale": ("FLOAT", {"default": 150.0, "min": 0.0, "max": 1000.0, "step": 1.0}),
|
|
"median_filter_size": ("INT", {"default": 31, "min": 1, "max": 31, "step": 2}),
|
|
"dithering_strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5.0, "step": 0.1}),
|
|
"spacer_1": ("STRING", {"multiline": False, "default": ""}),
|
|
"mask_enabled": ("BOOLEAN", {"default": True}),
|
|
"brush_intensity": ("FLOAT", {"default": 0.70, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"brush_softness": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 20.0, "step": 0.1}),
|
|
"mask_intensity_mult": ("FLOAT", {"default": 0.20, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"mask_gamma_mult": ("FLOAT", {"default": 0.14, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"spacer_2": ("STRING", {"multiline": False, "default": ""}),
|
|
"light_type_1": (s.LIGHT_TYPES, {"default": "Point"}),
|
|
"light_x_1": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"light_y_1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"light_z_1": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"intensity_1": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 50.0, "step": 0.1}),
|
|
"point_size_1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"neon_length_1": ("FLOAT", {"default": 0.77, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"neon_angle_1": ("FLOAT", {"default": 90.0, "min": 0.0, "max": 360.0, "step": 1.0}),
|
|
"ui_anchor": ("STRING", {"multiline": True, "default": ""}),
|
|
"enable_light_2": ("BOOLEAN", {"default": False}),
|
|
"light_type_2": (s.LIGHT_TYPES, {"default": "Point"}),
|
|
"light_x_2": ("FLOAT", {"default": 0.65, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"light_y_2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"light_z_2": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"intensity_2": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 50.0, "step": 0.1}),
|
|
"point_size_2": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"neon_length_2": ("FLOAT", {"default": 0.77, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"neon_angle_2": ("FLOAT", {"default": 90.0, "min": 0.0, "max": 360.0, "step": 1.0}),
|
|
},
|
|
"optional": { "mask": ("MASK",) }
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "relight_image"
|
|
CATEGORY = "Image/Lighting"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, **kwargs):
|
|
return float("NaN")
|
|
|
|
def parse_hex_color(self, hex_color: str) -> np.ndarray:
|
|
hex_color = hex_color.lstrip('#');
|
|
if len(hex_color) != 6: return np.array([1.0, 1.0, 1.0], dtype=np.float32)
|
|
try:
|
|
r = int(hex_color[0:2], 16) / 255.0; g = int(hex_color[2:4], 16) / 255.0; b = int(hex_color[4:6], 16) / 255.0
|
|
return np.array([r, g, b], dtype=np.float32)
|
|
except ValueError: return np.array([1.0, 1.0, 1.0], dtype=np.float32)
|
|
|
|
def tensor_to_numpy(self, tensor: torch.Tensor) -> np.ndarray: return tensor.squeeze(0).cpu().numpy()
|
|
|
|
def numpy_to_tensor(self, array: np.ndarray) -> torch.Tensor:
|
|
if array.ndim == 2: array = np.expand_dims(array, axis=-1)
|
|
if array.shape[2] == 1: array = np.repeat(array, 3, axis=-1)
|
|
return torch.from_numpy(array.astype(np.float32)).unsqueeze(0)
|
|
|
|
def tensor_to_pil(self, tensor):
|
|
return Image.fromarray(np.clip(255. * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def relight_image(self, image: torch.Tensor, normal_map: torch.Tensor, depth_map: torch.Tensor,
|
|
debug_view: str, preserve_color: bool, ambient_light: float, gamma: float, light_color: str,
|
|
depth_scale: float, median_filter_size: int, dithering_strength: float, spacer_1: str,
|
|
mask_enabled: bool, brush_intensity: float, brush_softness: float, mask_intensity_mult: float, mask_gamma_mult: float, spacer_2: str,
|
|
|
|
light_type_1: str, light_x_1: float, light_y_1: float, light_z_1: float, intensity_1: float, point_size_1: float, neon_length_1: float, neon_angle_1: float,
|
|
ui_anchor: str,
|
|
enable_light_2: bool, light_type_2: str, light_x_2: float, light_y_2: float, light_z_2: float, intensity_2: float, point_size_2: float, neon_length_2: float, neon_angle_2: float,
|
|
mask: torch.Tensor = None):
|
|
|
|
start_preview_server()
|
|
preview_image_pil = self.tensor_to_pil(image)
|
|
preview_filename = "relightpreview.png"
|
|
preview_filepath = os.path.join(folder_paths.get_temp_directory(), preview_filename)
|
|
preview_image_pil.save(preview_filepath)
|
|
full_url = f"http://127.0.0.1:{PREVIEW_SERVER_PORT}/{preview_filename}"
|
|
|
|
# VOTRE LOGIQUE DE TRAITEMENT ORIGINALE ET COMPLÈTE
|
|
img_np = self.tensor_to_numpy(image).astype(np.float32); height, width, _ = img_np.shape
|
|
normal_np = self.tensor_to_numpy(normal_map).astype(np.float32)
|
|
depth_np = self.tensor_to_numpy(depth_map).astype(np.float32)
|
|
np.nan_to_num(normal_np, copy=False, nan=0.0); np.nan_to_num(depth_np, copy=False, nan=0.0)
|
|
if normal_np.shape[:2] != (height, width): normal_np = cv2.resize(normal_np, (width, height), interpolation=cv2.INTER_LINEAR)
|
|
if len(depth_np.shape) < 3 or depth_np.shape[:2] != (height, width): depth_np = cv2.resize(depth_np, (width, height), interpolation=cv2.INTER_LINEAR)
|
|
if len(depth_np.shape) == 2: depth_np = np.expand_dims(depth_np, axis=-1)
|
|
normal_vectors = (normal_np * 2.0) - 1.0
|
|
xx, yy = np.meshgrid(np.arange(width), np.arange(height), indexing='xy')
|
|
pixel_z = (1.0 - depth_np[..., 0]) * depth_scale
|
|
pixel_positions = np.stack((xx.astype(np.float32), yy.astype(np.float32), pixel_z.astype(np.float32)), axis=-1)
|
|
light_vectors_1 = np.zeros_like(pixel_positions)
|
|
if light_type_1 == "Point":
|
|
light_pos_1 = np.array([light_x_1 * width, light_y_1 * height, light_z_1 * max(width, height)])
|
|
light_vectors_1 = light_pos_1 - pixel_positions
|
|
elif light_type_1 == "Neon":
|
|
center_x_1, center_y_1 = light_x_1 * width, light_y_1 * height; light_pos_z_1 = light_z_1 * max(width, height); angle_rad_1 = np.deg2rad(neon_angle_1); half_len_1 = (neon_length_1 / 2.0) * max(width, height); dx_1, dy_1 = half_len_1 * np.cos(angle_rad_1), half_len_1 * np.sin(angle_rad_1); p1_1 = np.array([center_x_1 - dx_1, center_y_1 - dy_1, light_pos_z_1]); p2_1 = np.array([center_x_1 + dx_1, center_y_1 + dy_1, light_pos_z_1]); line_vec_1 = p2_1 - p1_1; line_len_sq_1 = np.sum(line_vec_1**2)
|
|
if line_len_sq_1 == 0: line_len_sq_1 = 1.0
|
|
pixel_to_p1_1 = pixel_positions - p1_1; t_1 = np.sum(pixel_to_p1_1 * line_vec_1, axis=2) / line_len_sq_1; t_1 = np.clip(t_1, 0.0, 1.0); closest_points_1 = p1_1 + np.expand_dims(t_1, axis=2) * line_vec_1; light_vectors_1 = closest_points_1 - pixel_positions
|
|
dist_1 = np.linalg.norm(light_vectors_1, axis=2, keepdims=True); dist_1[dist_1 == 0] = 1.0
|
|
normalized_light_vectors_1 = light_vectors_1 / dist_1
|
|
normalized_light_vectors_2 = np.zeros_like(pixel_positions)
|
|
if enable_light_2:
|
|
light_vectors_2 = np.zeros_like(pixel_positions)
|
|
if light_type_2 == "Point":
|
|
light_pos_2 = np.array([light_x_2 * width, light_y_2 * height, light_z_2 * max(width, height)]); light_vectors_2 = light_pos_2 - pixel_positions
|
|
elif light_type_2 == "Neon":
|
|
center_x_2, center_y_2 = light_x_2 * width, light_y_2 * height; light_pos_z_2 = light_z_2 * max(width, height); angle_rad_2 = np.deg2rad(neon_angle_2); half_len_2 = (neon_length_2 / 2.0) * max(width, height); dx_2, dy_2 = half_len_2 * np.cos(angle_rad_2), half_len_2 * np.sin(angle_rad_2); p1_2 = np.array([center_x_2 - dx_2, center_y_2 - dy_2, light_pos_z_2]); p2_2 = np.array([center_x_2 + dx_2, center_y_2 + dy_2, light_pos_z_2]); line_vec_2 = p2_2 - p1_2; line_len_sq_2 = np.sum(line_vec_2**2)
|
|
if line_len_sq_2 == 0: line_len_sq_2 = 1.0
|
|
pixel_to_p1_2 = pixel_positions - p1_2; t_2 = np.sum(pixel_to_p1_2 * line_vec_2, axis=2) / line_len_sq_2; t_2 = np.clip(t_2, 0.0, 1.0); closest_points_2 = p1_2 + np.expand_dims(t_2, axis=2) * line_vec_2; light_vectors_2 = closest_points_2 - pixel_positions
|
|
dist_2 = np.linalg.norm(light_vectors_2, axis=2, keepdims=True); dist_2[dist_2 == 0] = 1.0
|
|
normalized_light_vectors_2 = light_vectors_2 / dist_2
|
|
dot_product_1 = np.sum(normal_vectors * normalized_light_vectors_1, axis=2)
|
|
if light_type_1 == "Point" and point_size_1 > 0: diffuse_term_1 = (dot_product_1 + point_size_1) / (1 + point_size_1)
|
|
else: diffuse_term_1 = dot_product_1
|
|
raw_light_1 = np.maximum(0, diffuse_term_1) * intensity_1
|
|
raw_light_2 = np.zeros_like(raw_light_1)
|
|
if enable_light_2:
|
|
dot_product_2 = np.sum(normal_vectors * normalized_light_vectors_2, axis=2)
|
|
if light_type_2 == "Point" and point_size_2 > 0: diffuse_term_2 = (dot_product_2 + point_size_2) / (1 + point_size_2)
|
|
else: diffuse_term_2 = dot_product_2
|
|
raw_light_2 = np.maximum(0, diffuse_term_2) * intensity_2
|
|
base_raw_light = raw_light_1 + raw_light_2
|
|
np.clip(base_raw_light, 0.0, 150.0, out=base_raw_light)
|
|
base_compressed_light = np.zeros_like(base_raw_light, dtype=np.float32)
|
|
if np.any(base_raw_light > 0): base_compressed_light[base_raw_light > 0] = np.power(base_raw_light[base_raw_light > 0], 1.0 / gamma)
|
|
final_lighting_map = base_compressed_light
|
|
shape_normals_vis = np.zeros((height, width, 3), dtype=np.float32); lighting_map_shape_vis = np.zeros((height, width), dtype=np.float32)
|
|
if mask is not None and mask_enabled:
|
|
mask_np = self.tensor_to_numpy(mask).squeeze();
|
|
if mask_np.shape[:2] != (height, width):
|
|
mask_np = cv2.resize(mask_np, (width, height), interpolation=cv2.INTER_LINEAR)
|
|
mask_np_binary = (mask_np > 0.5).astype(np.uint8)
|
|
shape_normals = np.zeros((height, width, 3), dtype=np.float32)
|
|
if np.any(mask_np_binary > 0):
|
|
if mask_np_binary.shape[:2] != (height, width): mask_np_binary = cv2.resize(mask_np_binary, (width, height), interpolation=cv2.INTER_NEAREST)
|
|
mask_bool = mask_np_binary > 0; dist_transform = distance_transform_edt(mask_np_binary); max_dist = np.max(dist_transform); normalized_dist = dist_transform / max_dist if max_dist > 0 else 0; grad_y, grad_x = np.gradient(normalized_dist)
|
|
shape_normals[mask_bool, 0] = -grad_x[mask_bool]; shape_normals[mask_bool, 1] = -grad_y[mask_bool]; shape_normals[mask_bool, 2] = 1.0
|
|
norm_sn = np.linalg.norm(shape_normals, axis=2, keepdims=True); norm_sn[norm_sn == 0] = 1; shape_normals /= norm_sn
|
|
shape_normals_vis = (shape_normals + 1.0) / 2.0
|
|
mask_dot_product_1 = np.sum(shape_normals * normalized_light_vectors_1, axis=2)
|
|
if light_type_1 == "Point" and point_size_1 > 0: mask_diffuse_1 = (mask_dot_product_1 + point_size_1) / (1 + point_size_1)
|
|
else: mask_diffuse_1 = mask_dot_product_1
|
|
mask_raw_light_1 = np.maximum(0, mask_diffuse_1) * intensity_1 * mask_intensity_mult
|
|
mask_raw_light_2 = np.zeros_like(mask_raw_light_1)
|
|
if enable_light_2:
|
|
mask_dot_product_2 = np.sum(shape_normals * normalized_light_vectors_2, axis=2)
|
|
if light_type_2 == "Point" and point_size_2 > 0: mask_diffuse_2 = (mask_dot_product_2 + point_size_2) / (1 + point_size_2)
|
|
else: mask_diffuse_2 = mask_dot_product_2
|
|
mask_raw_light_2 = np.maximum(0, mask_diffuse_2) * intensity_2 * mask_intensity_mult
|
|
mask_total_raw_light = mask_raw_light_1 + mask_raw_light_2
|
|
np.clip(mask_total_raw_light, 0.0, 150.0, out=mask_total_raw_light)
|
|
mask_compressed_light = np.zeros_like(mask_total_raw_light, dtype=np.float32)
|
|
final_mask_gamma = gamma * mask_gamma_mult;
|
|
if final_mask_gamma < 0.1: final_mask_gamma = 0.1
|
|
if np.any(mask_total_raw_light > 0): mask_compressed_light[mask_total_raw_light > 0] = np.power(mask_total_raw_light[mask_total_raw_light > 0], 1.0 / final_mask_gamma)
|
|
lighting_map_shape_vis = mask_compressed_light
|
|
blur_sigma = brush_softness
|
|
if blur_sigma > 0:
|
|
k_size = int(blur_sigma * 4) + 1; k_size = k_size if k_size % 2 != 0 else k_size + 1; blend_mask = cv2.GaussianBlur(mask_np.astype(np.float32), (k_size, k_size), blur_sigma)
|
|
else: blend_mask = mask_np.astype(np.float32)
|
|
blend_mask *= brush_intensity
|
|
final_lighting_map = (base_compressed_light * (1.0 - blend_mask)) + (mask_compressed_light * blend_mask)
|
|
final_lighting_map_expanded = np.expand_dims(final_lighting_map, axis=-1)
|
|
if median_filter_size > 1:
|
|
final_lighting_map_expanded = np.nan_to_num(final_lighting_map_expanded, nan=0.0)
|
|
k_size = median_filter_size if median_filter_size % 2 != 0 else median_filter_size + 1; squeezed_map = np.squeeze(final_lighting_map_expanded); min_val, max_val = np.min(squeezed_map), np.max(squeezed_map)
|
|
if max_val > min_val:
|
|
normalized_map = (squeezed_map - min_val) / (max_val - min_val); map_8bit = (normalized_map * 255).astype(np.uint8); filtered_8bit = cv2.medianBlur(map_8bit, k_size); filtered_normalized = filtered_8bit.astype(np.float32) / 255.0; restored_map = filtered_normalized * (max_val - min_val) + min_val; final_lighting_map_expanded = np.expand_dims(restored_map, axis=-1)
|
|
final_map_with_ambient = np.squeeze(final_lighting_map_expanded) + ambient_light
|
|
light_color_rgb = self.parse_hex_color(light_color); is_white_light = np.all(light_color_rgb == 1.0)
|
|
output_np = np.zeros_like(img_np, dtype=np.float32)
|
|
if debug_view == "Final Image":
|
|
colored_lighting_map = np.expand_dims(final_map_with_ambient, axis=-1) * light_color_rgb
|
|
if preserve_color and is_white_light:
|
|
hls_image = cv2.cvtColor(img_np, cv2.COLOR_RGB2HLS); hls_image[:, :, 1] = np.clip(hls_image[:, :, 1] * final_map_with_ambient, 0.0, 1.0); output_np = cv2.cvtColor(hls_image, cv2.COLOR_HLS2RGB)
|
|
else: output_np = np.clip(img_np * colored_lighting_map, 0.0, 1.0)
|
|
elif debug_view == "Lighting Map Only": output_np = final_map_with_ambient
|
|
elif debug_view == "Mask Lighting Only": output_np = lighting_map_shape_vis
|
|
elif debug_view == "Mask Virtual Normals": output_np = shape_normals_vis
|
|
if dithering_strength > 0 and debug_view == "Final Image":
|
|
noise = np.random.normal(0.0, dithering_strength / 255.0, output_np.shape).astype(np.float32)
|
|
output_np += noise
|
|
output_np = np.nan_to_num(np.clip(output_np, 0.0, 1.0))
|
|
|
|
return {
|
|
"result": (self.numpy_to_tensor(output_np),),
|
|
"ui": { "previews": [full_url] }
|
|
}
|
|
|
|
NODE_CLASS_MAPPINGS = { "RelightNode": RelightNode }
|
|
NODE_DISPLAY_NAME_MAPPINGS = { "RelightNode": "Advanced Relight Node" } |