stable-diffusion-webui-dump.../scripts/dumpunetlib/colorizer.py

157 lines
5.6 KiB
Python

from typing import Callable
import colorsys
import numpy as np
import numpy.typing as npt
"""
colorize transform fn (W/B: -∞..+∞ -> 0..1, R/B: -∞..+∞ -> -1..1)
---------------------------------------------------------------------------------------------
White/Black Linear v -> {|clamp(v, min, max)|-X}/(Y-X)
where
X = 0 if sign(min) != sign(max)
MIN(|min|,|max|) otherwise
Y = MAX(|min|,|max|)
Sigmoid v -> |sigmoid(v+offset, gain)-0.5|*2
Red/Blue Linear v -> [{clamp(v, min, max)-min}/(max-min)-0.5]*2
Sigmoid v -> {sigmoid(v+offset, gain)-0.5}*2
"""
def create_convert_linear_abs(min_: float, max_: float):
assert min_ < max_
X = 0.0 if np.sign(min_) != np.sign(max_) else \
min(np.abs(min_), np.abs(max_))
Y = max(np.abs(min_), np.abs(max_))
#@np.vectorize
def fn(array: npt.NDArray[np.float32]):
vs = (np.abs(np.clip(array, min_, max_)) - X) / (Y - X)
#assert np.all(((0.0 <= vs) & (vs <= 1.0)) | np.isnan(vs))
return vs
return fn
def create_convert_linear(min: float, max: float):
assert min < max
#@np.vectorize
def fn(array: npt.NDArray[np.float32]):
vs = ((np.clip(array, min, max) - min) / (max - min) - 0.5) * 2.0
#assert np.all(((-1.0 <= vs) & (vs <= 1.0)) | np.isnan(vs))
return vs
return fn
def convert_linear_auto01(array: npt.NDArray[np.float32]):
min_, max_ = np.min(array), np.max(array)
if np.isnan(min_) or np.isnan(max_) or min_ == max_:
return array
else:
return (array - min_) / (max_ - min_)
def convert_linear_auto11(array: npt.NDArray[np.float32]):
min_, max_ = np.min(array), np.max(array)
if np.isnan(min_) or np.isnan(max_) or min_ == max_:
return array
else:
return ((array - min_) / (max_ - min_) - 0.5) * 2.0
def create_convert_sigmoid_abs(gain: float, offset: float):
#@np.vectorize
def fn(array: npt.NDArray[np.float32]):
vs = 1.0 / (1.0 + np.exp(-gain * (array+offset)))
vs = np.abs(vs - 0.5) * 2.0
#assert np.all(((0.0 <= vs) & (vs <= 1.0)) | np.isnan(vs))
return vs
return fn
def create_convert_sigmoid(gain: float, offset: float):
#@np.vectorize
def fn(array: npt.NDArray[np.float32]):
vs = 1.0 / (1.0 + np.exp(-gain * (array+offset)))
vs = (vs - 0.5) * 2.0
#assert np.all(((-1.0 <= vs) & (vs <= 1.0)) | np.isnan(vs))
return vs
return fn
class Colorizer:
# v -> v
convert: Callable[[npt.NDArray[np.float32]], npt.NDArray[np.float32]]|None
# v -> pixel
colorize: Callable[[npt.NDArray[np.float32]], npt.NDArray[np.float32]]|None
format: str
def __init__(
self,
value1: str,
value2: str,
rgb: tuple[str,str,str],
hsl: tuple[str,str,str],
trans: str,
linear_minmax: tuple[float,float],
sigmoid_gain_offset: tuple[float,float]
):
assert value1 in ["White/Black", "Red/Blue", "Custom"]
if value1 == "White/Black":
self.format = "L"
colorize = None
elif value1 == "Red/Blue":
self.format = "RGB"
colorize = create_colorizer(colorize_red_blue_v)
else:
assert value2 in ["RGB", "HSL"]
if value2 == "RGB":
r, g, b = rgb
fn = eval(f"lambda v: ( ({r}), ({g}), ({b}) )", { "__builtins__": np }, {})
fv = np.vectorize(fn, otypes=[np.float32, np.float32, np.float32])
self.format ="RGB"
else:
h, s, l = hsl
fn = eval(f"lambda v: HLS2RGB( ({h}), ({l}), ({s}) )", { "__builtins__": np, "HLS2RGB": colorsys.hls_to_rgb }, {})
fv = np.vectorize(fn, otypes=[np.float32, np.float32, np.float32])
self.format = "RGB"
colorize = create_colorizer(fv)
assert trans in ["Auto [0,1]", "Auto [-1,1]", "Linear", "Sigmoid"]
if trans == "Auto [0,1]":
convert = convert_linear_auto01
elif trans == "Auto [-1,1]":
convert = convert_linear_auto11
elif trans == "Linear":
if value1 == "White/Black":
convert = create_convert_linear_abs(*linear_minmax)
else:
convert = create_convert_linear(*linear_minmax)
else:
if value1 == "White/Black":
convert = create_convert_sigmoid_abs(*sigmoid_gain_offset)
else:
convert = create_convert_sigmoid(*sigmoid_gain_offset)
self.convert = convert
self.colorize = colorize
def __call__(self, array: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
if self.convert is not None:
array = self.convert(array)
if self.colorize is not None:
array = self.colorize(array)
#assert np.all(((0.0 <= array) & (array <= 1.0)) | np.isnan(array))
return np.clip(array * 256, 0.0, 255.0).astype(np.uint8)
def colorize_red_blue_(v: float) -> tuple[float, float, float]:
# v = -∞..+∞ を
# v < 0 のとき青 (0, 0, v)
# v > 0 のとき赤 (v ,0, 0)
# にする
L = (v if v > 0.0 else 0.0, 0.0, np.abs(v) if v < 0.0 else 0.0)
return L
colorize_red_blue_v = np.vectorize(colorize_red_blue_, otypes=[np.float32, np.float32, np.float32])
def create_colorizer(fn: np.vectorize):
def colorizer(array: npt.NDArray[np.float32]):
return np.dstack(fn(array))
return colorizer