Skip to content

Commit ba319e4

Browse files
committed
Refactor node loading, add WASLUT
1 parent cb1688c commit ba319e4

File tree

1 file changed

+162
-5
lines changed

1 file changed

+162
-5
lines changed

nodes/WASLUT.py

Lines changed: 162 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -371,6 +371,24 @@ def _font():
371371
except Exception:
372372
return None
373373

374+
@staticmethod
375+
def _big_font():
376+
# Try to load a larger truetype font for better legibility; fallback to default
377+
# Common fonts to try across platforms
378+
candidates = [
379+
"DejaVuSansMono.ttf",
380+
"DejaVuSans.ttf",
381+
"Arial.ttf",
382+
"/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",
383+
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
384+
]
385+
for path in candidates:
386+
try:
387+
return ImageFont.truetype(path, size=14)
388+
except Exception:
389+
continue
390+
return WaveformScope._font()
391+
374392
@staticmethod
375393
def make_waveform_gray(ch_gray: np.ndarray, out_h: int) -> np.ndarray:
376394
h, w = ch_gray.shape
@@ -426,7 +444,7 @@ def compose_parade(wfr: np.ndarray, wfg: np.ndarray, wfb: np.ndarray,
426444
r_stats: tuple[float, float, float, float, float],
427445
g_stats: tuple[float, float, float, float, float],
428446
b_stats: tuple[float, float, float, float, float],
429-
gap: int = 8, pad: int = 60, left_pad: int = 56) -> np.ndarray:
447+
gap: int = 8, pad: int = 72, left_pad: int = 56) -> np.ndarray:
430448
h, w = wfr.shape
431449
pr = (np.stack([wfr, np.zeros_like(wfr), np.zeros_like(wfr)], -1) * 255.0 + 0.5).astype(np.uint8)
432450
pg = (np.stack([np.zeros_like(wfg), wfg, np.zeros_like(wfg)], -1) * 255.0 + 0.5).astype(np.uint8)
@@ -452,10 +470,21 @@ def compose_parade(wfr: np.ndarray, wfg: np.ndarray, wfb: np.ndarray,
452470
g_txt = f"G min {g_stats[0]:.4f} max {g_stats[1]:.4f} mean {g_stats[2]:.4f} std {g_stats[3]:.4f} median {g_stats[4]:.4f}"
453471
b_txt = f"B min {b_stats[0]:.4f} max {b_stats[1]:.4f} mean {b_stats[2]:.4f} std {b_stats[3]:.4f} median {b_stats[4]:.4f}"
454472
d = ImageDraw.Draw(canvas)
473+
big_font = WaveformScope._big_font()
474+
# Estimate line height for spacing
475+
try:
476+
bbox = big_font.getbbox("Ag")
477+
line_h = (bbox[3] - bbox[1]) + 4
478+
except Exception:
479+
line_h = 18
455480
y0 = H + 6
456-
d.text((6, y0), r_txt, fill=(255, 64, 64), font=font, stroke_width=1, stroke_fill=(0, 0, 0))
457-
d.text((6, y0 + 16), g_txt, fill=(64, 255, 64), font=font, stroke_width=1, stroke_fill=(0, 0, 0))
458-
d.text((6, y0 + 32), b_txt, fill=(64, 128, 255), font=font, stroke_width=1, stroke_fill=(0, 0, 0))
481+
# X positions aligned under each channel
482+
x_r = left_pad + 0 * (w + gap) + 6
483+
x_g = left_pad + 1 * (w + gap) + 6
484+
x_b = left_pad + 2 * (w + gap) + 6
485+
d.text((x_r, y0), r_txt, fill=(255, 64, 64), font=big_font, stroke_width=1, stroke_fill=(0, 0, 0))
486+
d.text((x_g, y0), g_txt, fill=(64, 255, 64), font=big_font, stroke_width=1, stroke_fill=(0, 0, 0))
487+
d.text((x_b, y0), b_txt, fill=(64, 128, 255), font=big_font, stroke_width=1, stroke_fill=(0, 0, 0))
459488
return np.array(canvas, dtype=np.uint8)
460489

461490
# Load LUT
@@ -498,6 +527,8 @@ def IS_CHANGED(cls, **kwargs):
498527
return None
499528

500529
RETURN_TYPES = ("LUT",)
530+
RETURN_NAMES = ("lut",)
531+
501532
FUNCTION = "run"
502533
CATEGORY = "WAS/Color/LUT"
503534

@@ -660,7 +691,6 @@ def _hsv_to_rgb(h: np.ndarray, s: np.ndarray, v: np.ndarray) -> np.ndarray:
660691

661692
@staticmethod
662693
def blend_hsv(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
663-
# Convert to HSV, circularly lerp hue, linearly lerp s and v
664694
ha, sa, va = LUTBlender._rgb_to_hsv(a)
665695
hb, sb, vb = LUTBlender._rgb_to_hsv(b)
666696
tt = float(t)
@@ -671,6 +701,123 @@ def blend_hsv(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
671701
out = LUTBlender._hsv_to_rgb(h, s, v)
672702
return np.clip(out, 0.0, 1.0).astype(np.float32)
673703

704+
@staticmethod
705+
def _srgb_to_linear(x: np.ndarray) -> np.ndarray:
706+
x = x.astype(np.float32)
707+
return np.where(x <= 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4).astype(np.float32)
708+
709+
@staticmethod
710+
def _linear_to_srgb(x: np.ndarray) -> np.ndarray:
711+
x = x.astype(np.float32)
712+
return np.where(x <= 0.0031308, x * 12.92, 1.055 * (np.clip(x, 0.0, None) ** (1/2.4)) - 0.055).astype(np.float32)
713+
714+
@staticmethod
715+
def _rgb_linear_to_xyz(rgb: np.ndarray) -> np.ndarray:
716+
M = np.array([
717+
[0.4124564, 0.3575761, 0.1804375],
718+
[0.2126729, 0.7151522, 0.0721750],
719+
[0.0193339, 0.1191920, 0.9503041],
720+
], dtype=np.float32)
721+
return np.tensordot(rgb, M.T, axes=1).astype(np.float32)
722+
723+
@staticmethod
724+
def _xyz_to_rgb_linear(xyz: np.ndarray) -> np.ndarray:
725+
M = np.array([
726+
[ 3.2404542, -1.5371385, -0.4985314],
727+
[-0.9692660, 1.8760108, 0.0415560],
728+
[ 0.0556434, -0.2040259, 1.0572252],
729+
], dtype=np.float32)
730+
return np.tensordot(xyz, M.T, axes=1).astype(np.float32)
731+
732+
@staticmethod
733+
def _rgb_to_lab(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
734+
lin = LUTBlender._srgb_to_linear(rgb)
735+
xyz = LUTBlender._rgb_linear_to_xyz(lin)
736+
Xn, Yn, Zn = 0.95047, 1.0, 1.08883
737+
x = xyz[..., 0] / Xn
738+
y = xyz[..., 1] / Yn
739+
z = xyz[..., 2] / Zn
740+
e = (6/29) ** 3
741+
k = (29/6) ** 2 / 3
742+
f = lambda t: np.where(t > e, np.cbrt(t), k * t + 4/29)
743+
fx, fy, fz = f(x), f(y), f(z)
744+
L = 116 * fy - 16
745+
a = 500 * (fx - fy)
746+
b = 200 * (fy - fz)
747+
return L.astype(np.float32), a.astype(np.float32), b.astype(np.float32)
748+
749+
@staticmethod
750+
def _lab_to_rgb(L: np.ndarray, a: np.ndarray, b: np.ndarray) -> np.ndarray:
751+
fy = (L + 16.0) / 116.0
752+
fx = fy + (a / 500.0)
753+
fz = fy - (b / 200.0)
754+
e = (6/29)
755+
e3 = e ** 3
756+
k = 3 * (e ** 2)
757+
invf = lambda t: np.where(t > e, t ** 3, (t - 4/29) / k)
758+
Xn, Yn, Zn = 0.95047, 1.0, 1.08883
759+
x = invf(fx) * Xn
760+
y = invf(fy) * Yn
761+
z = invf(fz) * Zn
762+
xyz = np.stack([x, y, z], axis=-1).astype(np.float32)
763+
lin = LUTBlender._xyz_to_rgb_linear(xyz)
764+
rgb = LUTBlender._linear_to_srgb(lin)
765+
return np.clip(rgb, 0.0, 1.0).astype(np.float32)
766+
767+
@staticmethod
768+
def _rgb_to_oklab(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
769+
lin = LUTBlender._srgb_to_linear(rgb)
770+
M1 = np.array([
771+
[0.4122214708, 0.5363325363, 0.0514459929],
772+
[0.2119034982, 0.6806995451, 0.1073969566],
773+
[0.0883024619, 0.2817188376, 0.6299787005],
774+
], dtype=np.float32)
775+
lms = np.tensordot(lin, M1.T, axes=1).astype(np.float32)
776+
l_, m_, s_ = np.cbrt(lms[..., 0]), np.cbrt(lms[..., 1]), np.cbrt(lms[..., 2])
777+
L = 0.2104542553 * l_ + 0.7936177850 * m_ - 0.0040720468 * s_
778+
a = 1.9779984951 * l_ - 2.4285922050 * m_ + 0.4505937099 * s_
779+
b = 0.0259040371 * l_ + 0.7827717662 * m_ - 0.8086757660 * s_
780+
return L.astype(np.float32), a.astype(np.float32), b.astype(np.float32)
781+
782+
@staticmethod
783+
def _oklab_to_rgb(L: np.ndarray, a: np.ndarray, b: np.ndarray) -> np.ndarray:
784+
l_ = L + 0.3963377774 * a + 0.2158037573 * b
785+
m_ = L - 0.1055613458 * a - 0.0638541728 * b
786+
s_ = L - 0.0894841775 * a - 1.2914855480 * b
787+
l = l_ ** 3
788+
m = m_ ** 3
789+
s = s_ ** 3
790+
M2 = np.array([
791+
[ 4.0767416621, -3.3077115913, 0.2309699292],
792+
[-1.2684380046, 2.6097574011, -0.3413193965],
793+
[-0.0041960863, -0.7034186147, 1.7076147010],
794+
], dtype=np.float32)
795+
lin = np.tensordot(np.stack([l, m, s], axis=-1), M2.T, axes=1).astype(np.float32)
796+
rgb = LUTBlender._linear_to_srgb(lin)
797+
return np.clip(rgb, 0.0, 1.0).astype(np.float32)
798+
799+
@staticmethod
800+
def blend_lab(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
801+
La, aa, ba = LUTBlender._rgb_to_lab(a)
802+
Lb, ab, bb = LUTBlender._rgb_to_lab(b)
803+
tt = float(t)
804+
L = La * (1.0 - tt) + Lb * tt
805+
A = aa * (1.0 - tt) + ab * tt
806+
B = ba * (1.0 - tt) + bb * tt
807+
out = LUTBlender._lab_to_rgb(L, A, B)
808+
return np.clip(out, 0.0, 1.0).astype(np.float32)
809+
810+
@staticmethod
811+
def blend_oklab(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
812+
La, aa, ba = LUTBlender._rgb_to_oklab(a)
813+
Lb, ab, bb = LUTBlender._rgb_to_oklab(b)
814+
tt = float(t)
815+
L = La * (1.0 - tt) + Lb * tt
816+
A = aa * (1.0 - tt) + ab * tt
817+
B = ba * (1.0 - tt) + bb * tt
818+
out = LUTBlender._oklab_to_rgb(L, A, B)
819+
return np.clip(out, 0.0, 1.0).astype(np.float32)
820+
674821
@staticmethod
675822
def blend_auto(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
676823
"""
@@ -699,6 +846,8 @@ def get_modes() -> list[str]:
699846
"smoothstep",
700847
"slerp",
701848
"hsv",
849+
"lab",
850+
"oklab",
702851
"auto",
703852
"multiply",
704853
"screen",
@@ -720,6 +869,7 @@ def INPUT_TYPES(cls):
720869
}
721870

722871
RETURN_TYPES = ("LUT",)
872+
RETURN_NAMES = ("lut",)
723873

724874
FUNCTION = "run"
725875
CATEGORY = "WAS/Color/LUT"
@@ -737,6 +887,10 @@ def run(self, lut_a, lut_b, mode, strength, output_size):
737887
C = LUTBlender.blend_slerp(A, B, strength)
738888
elif mode == "hsv":
739889
C = LUTBlender.blend_hsv(A, B, strength)
890+
elif mode == "lab":
891+
C = LUTBlender.blend_lab(A, B, strength)
892+
elif mode == "oklab":
893+
C = LUTBlender.blend_oklab(A, B, strength)
740894
elif mode == "auto":
741895
C = LUTBlender.blend_auto(A, B, strength)
742896
elif mode == "multiply":
@@ -762,6 +916,7 @@ def INPUT_TYPES(cls):
762916
}
763917

764918
RETURN_TYPES = ("IMAGE",)
919+
RETURN_NAMES = ("image",)
765920

766921
FUNCTION = "run"
767922
CATEGORY = "WAS/Color/LUT"
@@ -815,6 +970,7 @@ def INPUT_TYPES(cls):
815970
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE")
816971
RETURN_NAMES = ("red_waveform", "green_waveform", "blue_waveform", "rgb_parade")
817972
OUTPUT_NODE = True
973+
818974
FUNCTION = "run"
819975
CATEGORY = "WAS/Image/Scopes"
820976

@@ -861,6 +1017,7 @@ def run(self, image, waveform_height):
8611017

8621018
return {"ui": {"images": ui_entries}, "result": (red_batch, green_batch, blue_batch, parade_batch)}
8631019

1020+
8641021
NODE_CLASS_MAPPINGS = {
8651022
"WASLoadLUT": WASLoadLUT,
8661023
"WASCombineLUT": WASCombineLUT,

0 commit comments

Comments
 (0)