209 lines
7.4 KiB
Python
209 lines
7.4 KiB
Python
# Orion4D_SharpnessMax/utils.py
|
|
# Shared helpers for sharpness filters.
|
|
import math
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
EPS = 1e-6
|
|
|
|
# ---------- Conversions / clamps ----------
|
|
def ensure_tensor_bchw(img: torch.Tensor) -> torch.Tensor:
|
|
# ComfyUI images use [B,H,W,C] float32 [0..1].
|
|
if not isinstance(img, torch.Tensor):
|
|
raise TypeError("Expected torch.Tensor image")
|
|
if img.ndim != 4:
|
|
raise ValueError("Expected image with shape [B,H,W,C]")
|
|
if img.shape[-1] in (1, 3, 4): # B,H,W,C -> B,C,H,W
|
|
img = img.permute(0, 3, 1, 2).contiguous()
|
|
return img
|
|
|
|
def to_bhwc(img_bchw: torch.Tensor) -> torch.Tensor:
|
|
# B,C,H,W -> B,H,W,C
|
|
return img_bchw.permute(0, 2, 3, 1).contiguous()
|
|
|
|
def clamp01(x: torch.Tensor) -> torch.Tensor:
|
|
return x.clamp(0.0, 1.0)
|
|
|
|
# ---------- Gaussian blur ----------
|
|
def get_sigma_from_radius(radius: float) -> float:
|
|
# Soft Photoshop-like heuristic
|
|
return max(0.15, radius * 0.5 + 0.25)
|
|
|
|
def gaussian_kernel1d(sigma: float, truncate: float = 3.0,
|
|
device=None, dtype=None) -> torch.Tensor:
|
|
if sigma <= 0:
|
|
k = torch.tensor([1.0], device=device, dtype=dtype or torch.float32)
|
|
return k / k.sum()
|
|
radius = int(truncate * sigma + 0.5)
|
|
x = torch.arange(-radius, radius + 1, device=device, dtype=dtype or torch.float32)
|
|
kernel = torch.exp(-(x ** 2) / (2 * sigma * sigma))
|
|
kernel /= kernel.sum()
|
|
return kernel
|
|
|
|
def gaussian_blur(img_bchw: torch.Tensor, radius: float) -> torch.Tensor:
|
|
if radius <= 0:
|
|
return img_bchw
|
|
sigma = get_sigma_from_radius(radius)
|
|
device, dtype = img_bchw.device, img_bchw.dtype
|
|
k1d = gaussian_kernel1d(sigma, device=device, dtype=dtype)
|
|
|
|
kH = k1d.view(1, 1, -1, 1)
|
|
kW = k1d.view(1, 1, 1, -1)
|
|
C = img_bchw.shape[1]
|
|
|
|
pad_h = (kH.shape[2] - 1) // 2
|
|
pad_w = (kW.shape[3] - 1) // 2
|
|
|
|
x = F.pad(img_bchw, (0, 0, pad_h, pad_h), mode='reflect')
|
|
x = F.conv2d(x, kH.expand(C, 1, -1, 1), groups=C)
|
|
x = F.pad(x, (pad_w, pad_w, 0, 0), mode='reflect')
|
|
x = F.conv2d(x, kW.expand(C, 1, 1, -1), groups=C)
|
|
return x
|
|
|
|
# ---------- Luma / recomposition ----------
|
|
def rgb_to_luma(rgb_bchw: torch.Tensor) -> torch.Tensor:
|
|
# BT.709
|
|
r = rgb_bchw[:, 0:1]
|
|
g = rgb_bchw[:, 1:2]
|
|
b = rgb_bchw[:, 2:3]
|
|
return 0.2126 * r + 0.7152 * g + 0.0722 * b
|
|
|
|
def luma_to_rgb(luma_bchw: torch.Tensor, rgb_base_bchw: torch.Tensor) -> torch.Tensor:
|
|
# Replace luma while preserving chroma with a simple scale factor
|
|
baseL = rgb_to_luma(rgb_base_bchw)
|
|
scale = (luma_bchw + EPS) / (baseL + EPS)
|
|
return clamp01(rgb_base_bchw * scale)
|
|
|
|
# ---------- Simplified blend modes ----------
|
|
def blend_overlay(base: torch.Tensor, blend: torch.Tensor) -> torch.Tensor:
|
|
res = torch.where(base <= 0.5, 2 * base * blend, 1 - 2 * (1 - base) * (1 - blend))
|
|
return clamp01(res)
|
|
|
|
def blend_soft_light(base: torch.Tensor, blend: torch.Tensor) -> torch.Tensor:
|
|
# Approx W3C
|
|
res = (1 - 2 * blend) * base * base + 2 * blend * base
|
|
return clamp01(res)
|
|
|
|
def apply_opacity(base: torch.Tensor, effect: torch.Tensor, opacity: float) -> torch.Tensor:
|
|
return clamp01(base * (1 - opacity) + effect * opacity)
|
|
|
|
# ---------- Edge detection ----------
|
|
def sobel_edges(img_bchw: torch.Tensor) -> torch.Tensor:
|
|
# Sobel magnitude on luma, normalized to [0..1]
|
|
gray = rgb_to_luma(img_bchw)
|
|
kx = torch.tensor(
|
|
[[-1, 0, 1],
|
|
[-2, 0, 2],
|
|
[-1, 0, 1]], dtype=img_bchw.dtype, device=img_bchw.device
|
|
).view(1, 1, 3, 3)
|
|
ky = torch.tensor(
|
|
[[-1, -2, -1],
|
|
[ 0, 0, 0],
|
|
[ 1, 2, 1]], dtype=img_bchw.dtype, device=img_bchw.device
|
|
).view(1, 1, 3, 3)
|
|
gx = F.conv2d(F.pad(gray, (1, 1, 1, 1), mode='reflect'), kx)
|
|
gy = F.conv2d(F.pad(gray, (1, 1, 1, 1), mode='reflect'), ky)
|
|
mag = torch.sqrt(gx * gx + gy * gy)
|
|
mag = mag / (mag.amax(dim=[2, 3], keepdim=True) + EPS)
|
|
return mag # B,1,H,W
|
|
|
|
# ---------- Approximate edge-aware guided filter ----------
|
|
def guided_smooth(img_bchw: torch.Tensor, radius: float, eps: float = 1e-4) -> torch.Tensor:
|
|
"""
|
|
Lightweight edge-aware smoothing (approximate guided filter) applied to luma.
|
|
Important patch: force an ODD kernel to preserve the exact H, W shape.
|
|
"""
|
|
if radius <= 0:
|
|
return rgb_to_luma(img_bchw)
|
|
|
|
k = int(max(1, round(radius)))
|
|
if k % 2 == 0:
|
|
k += 1 # odd kernel is required
|
|
|
|
device, dtype = img_bchw.device, img_bchw.dtype
|
|
kernel = torch.ones(1, 1, k, k, device=device, dtype=dtype) / (k * k)
|
|
pad = (k // 2, k // 2, k // 2, k // 2) # L,R,T,B
|
|
|
|
L = rgb_to_luma(img_bchw)
|
|
|
|
mean_L = F.conv2d(F.pad(L, pad, mode="reflect"), kernel)
|
|
mean_L2 = F.conv2d(F.pad(L*L, pad, mode="reflect"), kernel)
|
|
var_L = mean_L2 - mean_L * mean_L
|
|
|
|
a = var_L / (var_L + eps)
|
|
b = (1 - a) * mean_L
|
|
|
|
mean_a = F.conv2d(F.pad(a, pad, mode="reflect"), kernel)
|
|
mean_b = F.conv2d(F.pad(b, pad, mode="reflect"), kernel)
|
|
|
|
return mean_a * L + mean_b # B,1,H,W
|
|
|
|
# ---------- Midtone mask for Clarity ----------
|
|
def midtone_mask(img_bchw: torch.Tensor, softness: float = 0.35) -> torch.Tensor:
|
|
"""
|
|
Mask near 1 around midtones (luma≈0.5), 0 toward shadows/highlights.
|
|
softness ∈ (0, 0.5] controls the plateau width.
|
|
"""
|
|
L = rgb_to_luma(img_bchw) # B,1,H,W
|
|
s = max(1e-6, min(0.5, float(softness)))
|
|
m = torch.cos((L - 0.5) * math.pi / (2 * s)).clamp(-1, 1)
|
|
m = (m + 1) / 2.0
|
|
m = torch.where((L > 0.5 + s) | (L < 0.5 - s), torch.zeros_like(m), m)
|
|
return m # B,1,H,W
|
|
|
|
|
|
# ---------- ComfyUI masks ----------
|
|
def ensure_mask_b1hw(mask: torch.Tensor, target_h: int = None, target_w: int = None, device=None, dtype=None) -> torch.Tensor:
|
|
"""Convert a ComfyUI MASK to [B,1,H,W] float [0..1].
|
|
|
|
Accepted formats: [H,W], [B,H,W], [B,H,W,1], [B,1,H,W].
|
|
The mask is resized when needed to match the image.
|
|
"""
|
|
if mask is None:
|
|
return None
|
|
if not isinstance(mask, torch.Tensor):
|
|
raise TypeError("Expected torch.Tensor mask")
|
|
|
|
m = mask
|
|
if device is not None:
|
|
m = m.to(device)
|
|
if dtype is not None:
|
|
m = m.to(dtype)
|
|
else:
|
|
m = m.to(torch.float32)
|
|
|
|
if m.ndim == 2:
|
|
m = m.unsqueeze(0).unsqueeze(0) # 1,1,H,W
|
|
elif m.ndim == 3:
|
|
m = m.unsqueeze(1) # B,1,H,W
|
|
elif m.ndim == 4:
|
|
if m.shape[-1] == 1: # B,H,W,1 -> B,1,H,W
|
|
m = m.permute(0, 3, 1, 2).contiguous()
|
|
elif m.shape[1] != 1:
|
|
m = m[:, :1, :, :]
|
|
else:
|
|
raise ValueError("Expected mask with shape [H,W], [B,H,W], [B,H,W,1] or [B,1,H,W]")
|
|
|
|
m = m.clamp(0.0, 1.0)
|
|
|
|
if target_h is not None and target_w is not None and (m.shape[-2] != target_h or m.shape[-1] != target_w):
|
|
m = F.interpolate(m, size=(target_h, target_w), mode="bilinear", align_corners=False).clamp(0.0, 1.0)
|
|
|
|
return m
|
|
|
|
|
|
def blur_mask(mask_b1hw: torch.Tensor, radius: float) -> torch.Tensor:
|
|
if mask_b1hw is None:
|
|
return None
|
|
return gaussian_blur(mask_b1hw, radius).clamp(0.0, 1.0)
|
|
|
|
|
|
def apply_mask_to_effect(base_bchw: torch.Tensor, effect_bchw: torch.Tensor, mask_b1hw: torch.Tensor = None,
|
|
mask_opacity: float = 1.0) -> torch.Tensor:
|
|
"""Blend base/effect using a white=effect, black=base mask."""
|
|
if mask_b1hw is None:
|
|
return clamp01(effect_bchw)
|
|
strength = float(max(0.0, min(1.0, mask_opacity)))
|
|
m = (mask_b1hw * strength).clamp(0.0, 1.0)
|
|
return clamp01(base_bchw * (1.0 - m) + effect_bchw * m)
|