375 lines
14 KiB
Python
375 lines
14 KiB
Python
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
try:
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
except Exception: # pragma: no cover
|
|
web = None
|
|
PromptServer = None
|
|
|
|
|
|
ROOT_DIR = Path(__file__).resolve().parent
|
|
PRESETS_DIR = ROOT_DIR / "presets_shadows"
|
|
PRESETS_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
PRESET_KEYS = {
|
|
"canvas_width", "canvas_height", "scale_percent", "shadow_mode",
|
|
"shadow_spread", "shadow_blur", "shadow_offset_x", "shadow_offset_y",
|
|
"shadow_opacity", "shadow_color", "background_mode", "background_color",
|
|
}
|
|
|
|
|
|
def _safe_preset_name(value):
|
|
value = str(value or "").strip()
|
|
if not value:
|
|
return ""
|
|
for ch in '<>:"/\\|?*':
|
|
value = value.replace(ch, "_")
|
|
value = value.replace("..", "_")
|
|
return value[:120].strip().rstrip(".")
|
|
|
|
|
|
def _preset_path(name):
|
|
safe = _safe_preset_name(name)
|
|
if not safe:
|
|
raise ValueError("Preset name is empty")
|
|
return PRESETS_DIR / f"{safe}.json"
|
|
|
|
|
|
def _clean_preset_data(data):
|
|
if not isinstance(data, dict):
|
|
raise ValueError("Preset data must be a JSON object")
|
|
return {key: data[key] for key in PRESET_KEYS if key in data}
|
|
|
|
|
|
def _normalize_hex(value, fallback="#000000"):
|
|
value = str(value or "").strip()
|
|
if not value.startswith("#"):
|
|
value = "#" + value
|
|
raw = value[1:]
|
|
try:
|
|
if len(raw) == 3:
|
|
raw = "".join(ch * 2 for ch in raw)
|
|
if len(raw) != 6:
|
|
raise ValueError("invalid hex length")
|
|
int(raw, 16)
|
|
return f"#{raw.upper()}"
|
|
except Exception:
|
|
return fallback.upper()
|
|
|
|
|
|
def _hex_to_rgb01(value, fallback="#000000"):
|
|
value = _normalize_hex(value, fallback)
|
|
raw = value[1:]
|
|
return np.array([
|
|
int(raw[0:2], 16),
|
|
int(raw[2:4], 16),
|
|
int(raw[4:6], 16),
|
|
], dtype=np.float32) / 255.0
|
|
|
|
|
|
def _resize_image_and_alpha(rgb, alpha, scale_factor):
|
|
h, w = rgb.shape[:2]
|
|
new_w = max(1, int(round(w * scale_factor)))
|
|
new_h = max(1, int(round(h * scale_factor)))
|
|
|
|
img_t = torch.from_numpy(np.ascontiguousarray(rgb)).permute(2, 0, 1).unsqueeze(0).float()
|
|
alpha_t = torch.from_numpy(np.ascontiguousarray(alpha)).unsqueeze(0).unsqueeze(0).float()
|
|
|
|
img_t = F.interpolate(img_t, size=(new_h, new_w), mode="bilinear", align_corners=False)
|
|
alpha_t = F.interpolate(alpha_t, size=(new_h, new_w), mode="bilinear", align_corners=False)
|
|
|
|
rgb_resized = img_t.squeeze(0).permute(1, 2, 0).cpu().numpy().astype(np.float32)
|
|
alpha_resized = alpha_t.squeeze(0).squeeze(0).cpu().numpy().astype(np.float32)
|
|
return rgb_resized, np.clip(alpha_resized, 0.0, 1.0)
|
|
|
|
|
|
def _place_center(canvas_h, canvas_w, item, fill_value=0.0):
|
|
if item.ndim == 3:
|
|
out = np.full((canvas_h, canvas_w, item.shape[2]), fill_value, dtype=np.float32)
|
|
item_h, item_w = item.shape[:2]
|
|
else:
|
|
out = np.full((canvas_h, canvas_w), fill_value, dtype=np.float32)
|
|
item_h, item_w = item.shape
|
|
|
|
dst_x0 = max(0, (canvas_w - item_w) // 2)
|
|
dst_y0 = max(0, (canvas_h - item_h) // 2)
|
|
dst_x1 = min(canvas_w, dst_x0 + item_w)
|
|
dst_y1 = min(canvas_h, dst_y0 + item_h)
|
|
|
|
src_x0 = max(0, (item_w - canvas_w) // 2)
|
|
src_y0 = max(0, (item_h - canvas_h) // 2)
|
|
src_x1 = src_x0 + (dst_x1 - dst_x0)
|
|
src_y1 = src_y0 + (dst_y1 - dst_y0)
|
|
|
|
if item.ndim == 3:
|
|
out[dst_y0:dst_y1, dst_x0:dst_x1, :] = item[src_y0:src_y1, src_x0:src_x1, :]
|
|
else:
|
|
out[dst_y0:dst_y1, dst_x0:dst_x1] = item[src_y0:src_y1, src_x0:src_x1]
|
|
return out
|
|
|
|
|
|
def _shift_mask(mask, offset_x, offset_y):
|
|
h, w = mask.shape
|
|
out = np.zeros_like(mask)
|
|
|
|
src_x0 = max(0, -offset_x)
|
|
src_y0 = max(0, -offset_y)
|
|
src_x1 = min(w, w - offset_x) if offset_x >= 0 else w
|
|
src_y1 = min(h, h - offset_y) if offset_y >= 0 else h
|
|
|
|
dst_x0 = max(0, offset_x)
|
|
dst_y0 = max(0, offset_y)
|
|
dst_x1 = dst_x0 + max(0, src_x1 - src_x0)
|
|
dst_y1 = dst_y0 + max(0, src_y1 - src_y0)
|
|
|
|
if src_x1 > src_x0 and src_y1 > src_y0 and dst_x1 > dst_x0 and dst_y1 > dst_y0:
|
|
out[dst_y0:dst_y1, dst_x0:dst_x1] = mask[src_y0:src_y1, src_x0:src_x1]
|
|
return out
|
|
|
|
|
|
def _gaussian_blur(mask, blur_px):
|
|
blur_px = float(max(0.0, blur_px))
|
|
if blur_px <= 0.0:
|
|
return mask.astype(np.float32)
|
|
|
|
sigma = max(0.5, blur_px * 0.35)
|
|
radius = max(1, int(math.ceil(sigma * 3.0)))
|
|
x = torch.arange(-radius, radius + 1, dtype=torch.float32)
|
|
kernel = torch.exp(-(x * x) / (2.0 * sigma * sigma))
|
|
kernel = kernel / kernel.sum()
|
|
|
|
t = torch.from_numpy(mask.astype(np.float32)).unsqueeze(0).unsqueeze(0)
|
|
kx = kernel.view(1, 1, 1, -1)
|
|
ky = kernel.view(1, 1, -1, 1)
|
|
|
|
t = F.pad(t, (radius, radius, 0, 0), mode="constant", value=0.0)
|
|
t = F.conv2d(t, kx)
|
|
t = F.pad(t, (0, 0, radius, radius), mode="constant", value=0.0)
|
|
t = F.conv2d(t, ky)
|
|
return t.squeeze(0).squeeze(0).cpu().numpy().astype(np.float32)
|
|
|
|
|
|
def _dilate_mask(mask, spread_px):
|
|
spread_px = int(max(0, spread_px))
|
|
if spread_px <= 0:
|
|
return mask.astype(np.float32)
|
|
k = spread_px * 2 + 1
|
|
t = torch.from_numpy(mask.astype(np.float32)).unsqueeze(0).unsqueeze(0)
|
|
t = F.max_pool2d(t, kernel_size=k, stride=1, padding=spread_px)
|
|
return t.squeeze(0).squeeze(0).cpu().numpy().astype(np.float32)
|
|
|
|
|
|
def _composite_layers(bg_rgb, bg_alpha, shadow_rgb, shadow_alpha, obj_rgb, obj_alpha):
|
|
bg_p = bg_rgb * bg_alpha[..., None]
|
|
sh_p = shadow_rgb[None, None, :] * shadow_alpha[..., None]
|
|
under_p = sh_p + bg_p * (1.0 - shadow_alpha[..., None])
|
|
under_a = shadow_alpha + bg_alpha * (1.0 - shadow_alpha)
|
|
|
|
obj_p = obj_rgb * obj_alpha[..., None]
|
|
out_p = obj_p + under_p * (1.0 - obj_alpha[..., None])
|
|
out_a = obj_alpha + under_a * (1.0 - obj_alpha)
|
|
|
|
out_rgb = np.zeros_like(obj_rgb, dtype=np.float32)
|
|
mask = out_a > 1e-6
|
|
out_rgb[mask] = out_p[mask] / out_a[mask, None]
|
|
return np.clip(out_rgb, 0.0, 1.0), np.clip(out_a, 0.0, 1.0)
|
|
|
|
|
|
class Orion4D_DropShadow:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"canvas_width": ("INT", {"default": 1024, "min": 1, "max": 8192, "step": 1}),
|
|
"canvas_height": ("INT", {"default": 1024, "min": 1, "max": 8192, "step": 1}),
|
|
"scale_percent": ("FLOAT", {"default": -15.0, "min": -99.0, "max": 500.0, "step": 0.1}),
|
|
"shadow_mode": (["outer", "inner"], {"default": "outer"}),
|
|
"shadow_spread": ("INT", {"default": 0, "min": 0, "max": 256, "step": 1}),
|
|
"shadow_blur": ("FLOAT", {"default": 24.0, "min": 0.0, "max": 256.0, "step": 0.5}),
|
|
"shadow_offset_x": ("INT", {"default": 24, "min": -4096, "max": 4096, "step": 1}),
|
|
"shadow_offset_y": ("INT", {"default": 24, "min": -4096, "max": 4096, "step": 1}),
|
|
"shadow_opacity": ("FLOAT", {"default": 0.55, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"shadow_color": ("STRING", {"default": "#000000"}),
|
|
"background_mode": (["transparent", "color"], {"default": "transparent"}),
|
|
"background_color": ("STRING", {"default": "#FFFFFF"}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask", "rgba_image", "shadow_rgba", "shadow_mask")
|
|
FUNCTION = "apply_shadow"
|
|
CATEGORY = "Orion4D_Layer"
|
|
|
|
def apply_shadow(
|
|
self,
|
|
image,
|
|
canvas_width,
|
|
canvas_height,
|
|
scale_percent,
|
|
shadow_mode,
|
|
shadow_spread,
|
|
shadow_blur,
|
|
shadow_offset_x,
|
|
shadow_offset_y,
|
|
shadow_opacity,
|
|
shadow_color,
|
|
background_mode,
|
|
background_color,
|
|
mask=None,
|
|
):
|
|
canvas_w = max(1, int(canvas_width))
|
|
canvas_h = max(1, int(canvas_height))
|
|
scale_factor = max(0.01, 1.0 + float(scale_percent) / 100.0)
|
|
sh_color = _hex_to_rgb01(shadow_color, "#000000")
|
|
bg_color = _hex_to_rgb01(background_color, "#FFFFFF")
|
|
bg_mode = background_mode if background_mode in ["transparent", "color"] else "transparent"
|
|
sh_mode = shadow_mode if shadow_mode in ["outer", "inner"] else "outer"
|
|
shadow_opacity = float(max(0.0, min(float(shadow_opacity), 1.0)))
|
|
spread = int(max(0, shadow_spread))
|
|
|
|
batch = image.shape[0]
|
|
out_images = []
|
|
out_masks = []
|
|
out_rgba = []
|
|
out_shadow_rgba = []
|
|
out_shadow_masks = []
|
|
|
|
mask_batch = None
|
|
if mask is not None:
|
|
mask_batch = mask
|
|
if len(mask_batch.shape) == 2:
|
|
mask_batch = mask_batch.unsqueeze(0)
|
|
|
|
for i in range(batch):
|
|
img = image[i].detach().cpu().numpy().astype(np.float32)
|
|
rgb = np.clip(img[..., :3], 0.0, 1.0)
|
|
|
|
if mask_batch is not None:
|
|
mi = min(i, mask_batch.shape[0] - 1)
|
|
alpha = 1.0 - np.clip(mask_batch[mi].detach().cpu().numpy().astype(np.float32), 0.0, 1.0)
|
|
elif img.shape[-1] >= 4:
|
|
alpha = np.clip(img[..., 3], 0.0, 1.0)
|
|
else:
|
|
alpha = np.ones(rgb.shape[:2], dtype=np.float32)
|
|
|
|
rgb_scaled, alpha_scaled = _resize_image_and_alpha(rgb, alpha, scale_factor)
|
|
rgb_canvas = _place_center(canvas_h, canvas_w, rgb_scaled, 0.0)
|
|
alpha_canvas = _place_center(canvas_h, canvas_w, alpha_scaled, 0.0)
|
|
|
|
if sh_mode == "outer":
|
|
shadow_seed = _dilate_mask(alpha_canvas, spread)
|
|
shadow_seed = _shift_mask(shadow_seed, int(shadow_offset_x), int(shadow_offset_y))
|
|
shadow_alpha = _gaussian_blur(shadow_seed, float(shadow_blur)) * shadow_opacity
|
|
shadow_alpha = np.clip(shadow_alpha, 0.0, 1.0)
|
|
else:
|
|
inner_seed = _dilate_mask(alpha_canvas, spread)
|
|
inner_seed = _shift_mask(inner_seed, int(shadow_offset_x), int(shadow_offset_y))
|
|
inner_seed = alpha_canvas * (1.0 - inner_seed)
|
|
shadow_alpha = _gaussian_blur(inner_seed, float(shadow_blur)) * shadow_opacity
|
|
shadow_alpha = np.clip(shadow_alpha * alpha_canvas, 0.0, 1.0)
|
|
|
|
# Full RGBA output: keep true transparency if requested.
|
|
if bg_mode == "color":
|
|
bg_alpha_rgba = np.ones((canvas_h, canvas_w), dtype=np.float32)
|
|
bg_rgb_rgba = np.broadcast_to(bg_color.reshape(1, 1, 3), (canvas_h, canvas_w, 3)).copy().astype(np.float32)
|
|
else:
|
|
bg_alpha_rgba = np.zeros((canvas_h, canvas_w), dtype=np.float32)
|
|
bg_rgb_rgba = np.zeros((canvas_h, canvas_w, 3), dtype=np.float32)
|
|
|
|
out_rgb_rgba, out_alpha = _composite_layers(bg_rgb_rgba, bg_alpha_rgba, sh_color, shadow_alpha, rgb_canvas, alpha_canvas)
|
|
rgba = np.concatenate([out_rgb_rgba, out_alpha[..., None]], axis=-1)
|
|
|
|
# Visible RGB output always uses the background color as backdrop.
|
|
bg_alpha_rgb = np.ones((canvas_h, canvas_w), dtype=np.float32)
|
|
bg_rgb_rgb = np.broadcast_to(bg_color.reshape(1, 1, 3), (canvas_h, canvas_w, 3)).copy().astype(np.float32)
|
|
out_rgb_visible, _ = _composite_layers(bg_rgb_rgb, bg_alpha_rgb, sh_color, shadow_alpha, rgb_canvas, alpha_canvas)
|
|
|
|
shadow_rgba = np.concatenate([
|
|
np.broadcast_to(sh_color.reshape(1, 1, 3), (canvas_h, canvas_w, 3)).copy().astype(np.float32),
|
|
shadow_alpha[..., None]
|
|
], axis=-1)
|
|
|
|
out_images.append(torch.from_numpy(np.ascontiguousarray(out_rgb_visible)).unsqueeze(0))
|
|
out_masks.append(torch.from_numpy(np.ascontiguousarray(out_alpha)).unsqueeze(0))
|
|
out_rgba.append(torch.from_numpy(np.ascontiguousarray(rgba)).unsqueeze(0))
|
|
out_shadow_rgba.append(torch.from_numpy(np.ascontiguousarray(shadow_rgba)).unsqueeze(0))
|
|
out_shadow_masks.append(torch.from_numpy(np.ascontiguousarray(shadow_alpha)).unsqueeze(0))
|
|
|
|
return (
|
|
torch.cat(out_images, dim=0),
|
|
torch.cat(out_masks, dim=0),
|
|
torch.cat(out_rgba, dim=0),
|
|
torch.cat(out_shadow_rgba, dim=0),
|
|
torch.cat(out_shadow_masks, dim=0),
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Orion4D_DropShadow": Orion4D_DropShadow,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Orion4D_DropShadow": "🌑 Drop Shadow",
|
|
}
|
|
|
|
|
|
if PromptServer is not None and web is not None:
|
|
@PromptServer.instance.routes.get("/orion4d/drop_shadow_presets")
|
|
async def orion4d_drop_shadow_presets_list(request):
|
|
presets = sorted(
|
|
(path.stem for path in PRESETS_DIR.glob("*.json")),
|
|
key=str.casefold,
|
|
)
|
|
return web.json_response({"presets": presets})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/orion4d/drop_shadow_presets/{name}")
|
|
async def orion4d_drop_shadow_preset_load(request):
|
|
try:
|
|
path = _preset_path(request.match_info.get("name", ""))
|
|
if not path.exists():
|
|
return web.json_response({"error": "Preset not found"}, status=404)
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
return web.json_response({"name": path.stem, "data": _clean_preset_data(data)})
|
|
except Exception as exc:
|
|
return web.json_response({"error": str(exc)}, status=400)
|
|
|
|
|
|
@PromptServer.instance.routes.post("/orion4d/drop_shadow_presets/save")
|
|
async def orion4d_drop_shadow_preset_save(request):
|
|
try:
|
|
payload = await request.json()
|
|
path = _preset_path(payload.get("name", ""))
|
|
data = _clean_preset_data(payload.get("data", {}))
|
|
path.write_text(
|
|
json.dumps(data, ensure_ascii=False, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
return web.json_response({"ok": True, "name": path.stem})
|
|
except Exception as exc:
|
|
return web.json_response({"error": str(exc)}, status=400)
|
|
|
|
|
|
@PromptServer.instance.routes.post("/orion4d/drop_shadow_presets/delete")
|
|
async def orion4d_drop_shadow_preset_delete(request):
|
|
try:
|
|
payload = await request.json()
|
|
path = _preset_path(payload.get("name", ""))
|
|
if path.exists():
|
|
path.unlink()
|
|
return web.json_response({"ok": True, "name": path.stem})
|
|
except Exception as exc:
|
|
return web.json_response({"error": str(exc)}, status=400)
|