diff --git a/Nodes/Rasterix.py b/Nodes/Rasterix.py index 75b145a..147b13f 100644 --- a/Nodes/Rasterix.py +++ b/Nodes/Rasterix.py @@ -22,6 +22,9 @@ from ..components.images import histogram as histogram from ..components.images import img_posterize as img_posterize from ..components.images import img_solarization_bw as img_solarization_bw from ..components.images import img_clarity as img_clarity +from ..components.images import img_dehaze as img_dehaze +from ..components.images import img_local_laplacian as img_local_laplacian +from ..components.images import img_frequency_separation as img_frequency_separation from ..components import utility from .Dashboard import PrimereModelConceptSelector as PrimereModelConceptSelector import os @@ -1001,4 +1004,87 @@ class PrimereClarity: if use_clarity and strength != 0: pil_img = img_clarity.img_clarity(image=pil_img, strength=strength, radius=radius, midtone_range=midtone_range, edge_preservation=edge_preservation, precision=precision) + return (utility.image_to_tensor(pil_img),) + +class PrimereDehaze: + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("IMAGE",) + FUNCTION = "primere_dehaze" + CATEGORY = TREE_RASTERIX + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {"forceInput": True}), + "use_dehaze": ("BOOLEAN", {"default": False, "label_off": "Ignore dehaze", "label_on": "Apply dehaze"}), + "precision": ("BOOLEAN", {"default": False, "label_off": "8 bit", "label_on": "16 bit"}), + + "strength": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}), + "radius": ("INT", {"default": 15, "min": 3, "max": 100, "step": 1}), + "omega": ("FLOAT", {"default": 0.95, "min": 0.5, "max": 1.0, "step": 0.01}), + "t0": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.5, "step": 0.01}), + "contrast": ("FLOAT", {"default": 1.05, "min": 0.5, "max": 2.0, "step": 0.01}), + } + } + + def primere_dehaze(self, image, use_dehaze, precision, strength, radius, omega, t0, contrast): + pil_img = utility.tensor_to_image(image) + if use_dehaze and strength > 0: + pil_img = img_dehaze.img_dehaze(image=pil_img, strength=strength, radius=radius, omega=omega, t0=t0, contrast=contrast, precision=precision) + + return (utility.image_to_tensor(pil_img),) + +class PrimereLocalLaplacian: + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("IMAGE",) + FUNCTION = "primere_local_laplacian" + CATEGORY = TREE_RASTERIX + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {"forceInput": True}), + "use_local_laplacian": ("BOOLEAN", {"default": False, "label_off": "Ignore local laplacian", "label_on": "Apply local laplacian"}), + + "sigma": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 5.0, "step": 0.1}), + "contrast": ("FLOAT", {"default": 1.2, "min": 0.5, "max": 3.0, "step": 0.01}), + "detail": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), + "levels": ("INT", {"default": 8, "min": 4, "max": 32, "step": 1}), + } + } + + def primere_local_laplacian(self, image, use_local_laplacian, sigma, contrast, detail, levels): + pil_img = utility.tensor_to_image(image) + if use_local_laplacian: + pil_img = img_local_laplacian.img_local_laplacian(image=pil_img, sigma=sigma, contrast=contrast, detail=detail, levels=levels) + + return (utility.image_to_tensor(pil_img),) + +class PrimereFrequencySeparation: + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("IMAGE",) + FUNCTION = "primere_frequency_separation" + CATEGORY = TREE_RASTERIX + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {"forceInput": True}), + "use_frequency_separation": ("BOOLEAN", {"default": False, "label_off": "Ignore frequency separation", "label_on": "Apply frequency separation"}), + + "radius": ("FLOAT", {"default": 3.0, "min": 0.5, "max": 20.0, "step": 0.1}), + "low_freq_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), + "high_freq_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), + "blend_mode": (["add", "multiply", "overlay"], {"default": "add"}), + } + } + + def primere_frequency_separation(self, image, use_frequency_separation, radius, low_freq_strength, high_freq_strength, blend_mode): + pil_img = utility.tensor_to_image(image) + if use_frequency_separation: + pil_img = img_frequency_separation.img_frequency_separation(image=pil_img, radius=radius, low_freq_strength=low_freq_strength, high_freq_strength=high_freq_strength, blend_mode=blend_mode) + return (utility.image_to_tensor(pil_img),) \ No newline at end of file diff --git a/__init__.py b/__init__.py index c9efed2..a3d57cd 100644 --- a/__init__.py +++ b/__init__.py @@ -80,6 +80,9 @@ NODE_CLASS_MAPPINGS = { "PrimereSolarizationBW": Rasterix.PrimereSolarizationBW, "PrimereClarity": Rasterix.PrimereClarity, "PrimereHistogram": Rasterix.PrimereHistogram, + "PrimereDehaze": Rasterix.PrimereDehaze, + "PrimereLocalLaplacian": Rasterix.PrimereLocalLaplacian, + "PrimereFrequencySeparation": Rasterix.PrimereFrequencySeparation, "PrimerePrompt": Inputs.PrimereDoublePrompt, "PrimereStyleLoader": Inputs.PrimereStyleLoader, @@ -180,6 +183,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PrimereSolarizationBW": "Primere Rasterix (Solarization)", "PrimereClarity": "Primere Rasterix (Clarity)", "PrimereHistogram": "Primere Rasterix (Histogram)", + "PrimereDehaze": "Primere Rasterix (Dehaze)", + "PrimereLocalLaplacian": "Primere Rasterix (Local Laplacian)", + "PrimereFrequencySeparation": "Primere Rasterix (Frequency Separation)", "PrimerePrompt": "Primere Prompt", "PrimereStyleLoader": "Primere Styles", diff --git a/components/images/img_dehaze.py b/components/images/img_dehaze.py new file mode 100644 index 0000000..2afab1f --- /dev/null +++ b/components/images/img_dehaze.py @@ -0,0 +1,44 @@ +import numpy as np +from PIL import Image +from scipy.ndimage import minimum_filter, gaussian_filter + + +def _estimate_atmospheric_light(arr: np.ndarray) -> np.ndarray: + flat = arr.reshape(-1, 3) + brightest = flat[np.argmax(np.sum(flat, axis=1))] + return brightest + + +def img_dehaze( + image: Image.Image, + strength: float = 0.7, + radius: int = 15, + omega: float = 0.95, + t0: float = 0.1, + contrast: float = 1.05, + precision: bool = False, +) -> Image.Image: + + img = image.convert("RGB") + + if precision: + max_val = 65535.0 + arr = np.array(img, dtype=np.float32) / 255.0 + arr = arr * max_val + arr = arr / max_val + else: + arr = np.array(img, dtype=np.float32) / 255.0 + + dark = np.min(arr, axis=2) + dark = minimum_filter(dark, size=radius) + A = _estimate_atmospheric_light(arr) + transmission = 1.0 - omega * dark + transmission = np.clip(transmission, t0, 1.0) + transmission = gaussian_filter(transmission, sigma=radius * 0.25) + J = (arr - A) / transmission[..., None] + A + out = arr * (1.0 - strength) + J * strength + out = (out - 0.5) * contrast + 0.5 + out = np.clip(out, 0.0, 1.0) + out = (out * 255.0).astype(np.uint8) + + return Image.fromarray(out, mode="RGB") \ No newline at end of file diff --git a/components/images/img_frequency_separation.py b/components/images/img_frequency_separation.py new file mode 100644 index 0000000..49ca690 --- /dev/null +++ b/components/images/img_frequency_separation.py @@ -0,0 +1,37 @@ +import numpy as np +from PIL import Image +from scipy.ndimage import gaussian_filter + + +def img_frequency_separation( + image: Image.Image, + radius: float = 3.0, + low_freq_strength: float = 1.0, + high_freq_strength: float = 1.0, + blend_mode: str = "add", +) -> Image.Image: + + img = image.convert("RGB") + arr = np.array(img, dtype=np.float32) / 255.0 + low = gaussian_filter(arr, sigma=(radius, radius, 0)) + high = arr - low + low_mod = low * low_freq_strength + high_mod = high * high_freq_strength + if blend_mode == "add": + out = low_mod + high_mod + elif blend_mode == "multiply": + out = low_mod * (1.0 + high_mod) + elif blend_mode == "overlay": + base = low_mod + detail = high_mod + out = np.where( + base <= 0.5, + 2.0 * base * (1.0 + detail), + 1.0 - 2.0 * (1.0 - base) * (1.0 - detail) + ) + else: + out = low_mod + high_mod + out = np.clip(out, 0.0, 1.0) + out = (out * 255.0).astype(np.uint8) + + return Image.fromarray(out, mode="RGB") \ No newline at end of file diff --git a/components/images/img_local_laplacian.py b/components/images/img_local_laplacian.py new file mode 100644 index 0000000..84bf78f --- /dev/null +++ b/components/images/img_local_laplacian.py @@ -0,0 +1,34 @@ +import numpy as np +from PIL import Image +from scipy.ndimage import gaussian_filter + + +def _to_luminance(arr: np.ndarray) -> np.ndarray: + return 0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2] + + +def img_local_laplacian( + image: Image.Image, + sigma: float = 1.0, + contrast: float = 1.2, + detail: float = 1.0, + levels: int = 8, +) -> Image.Image: + + img = image.convert("RGB") + arr = np.array(img, dtype=np.float32) / 255.0 + luma = _to_luminance(arr) + base = gaussian_filter(luma, sigma=sigma) + detail_layer = luma - base + remapped = base + contrast * (base - 0.5) + step = 1.0 / levels + quantized = np.floor(luma / step) * step + step * 0.5 + tone = (remapped * 0.7 + quantized * 0.3) + enhanced = tone + detail * detail_layer + enhanced = np.clip(enhanced, 0.0, 1.0) + scale = enhanced / (luma + 1e-6) + out = arr * scale[..., None] + out = np.clip(out, 0.0, 1.0) + out = (out * 255.0).astype(np.uint8) + + return Image.fromarray(out, mode="RGB") \ No newline at end of file