Upload models/oklab_utils.py with huggingface_hub

models/oklab_utils.py

@@ -0,0 +1,263 @@
+"""
+OKLab Color Space Utilities
+
+Perceptually uniform color space for semantic loss computation.
+OKLab ensures that equal distances in the color space correspond to
+equal perceived differences — critical for meaningful color-based encoding.
+
+Key functions:
+- srgb_to_oklab / oklab_to_srgb: Color space conversions
+- rotate_ab: Rotate hue in a-b plane (for domain/idiom shifts)
+- set_chroma: Set chroma magnitude (for purity encoding)
+- OKLabMSELoss: Perceptually uniform loss function
+- hsl_to_oklab_batch: Batch conversion for training
+"""
+
+import torch
+import torch.nn as nn
+import math
+from typing import Tuple
+
+
+def clamp(x: float, lo: float, hi: float) -> float:
+    """Clamp a value to [lo, hi]."""
+    return max(lo, min(hi, x))
+
+
+# ── sRGB ↔ Linear RGB ──
+
+def srgb_to_linear(c: float) -> float:
+    """sRGB gamma to linear."""
+    if c <= 0.04045:
+        return c / 12.92
+    return ((c + 0.055) / 1.055) ** 2.4
+
+
+def linear_to_srgb(c: float) -> float:
+    """Linear to sRGB gamma."""
+    if c <= 0.0031308:
+        return c * 12.92
+    return 1.055 * (c ** (1.0 / 2.4)) - 0.055
+
+
+# ── sRGB ↔ OKLab ──
+
+def srgb_to_oklab(r: float, g: float, b: float) -> Tuple[float, float, float]:
+    """Convert sRGB [0,1] to OKLab."""
+    r_lin = srgb_to_linear(r)
+    g_lin = srgb_to_linear(g)
+    b_lin = srgb_to_linear(b)
+
+    l_ = 0.4122214708 * r_lin + 0.5363325363 * g_lin + 0.0514459929 * b_lin
+    m_ = 0.2119034982 * r_lin + 0.6806995451 * g_lin + 0.1073969566 * b_lin
+    s_ = 0.0883024619 * r_lin + 0.2817188376 * g_lin + 0.6299787005 * b_lin
+
+    l_c = l_ ** (1.0 / 3.0) if l_ >= 0 else -((-l_) ** (1.0 / 3.0))
+    m_c = m_ ** (1.0 / 3.0) if m_ >= 0 else -((-m_) ** (1.0 / 3.0))
+    s_c = s_ ** (1.0 / 3.0) if s_ >= 0 else -((-s_) ** (1.0 / 3.0))
+
+    L = 0.2104542553 * l_c + 0.7936177850 * m_c - 0.0040720468 * s_c
+    a = 1.9779984951 * l_c - 2.4285922050 * m_c + 0.4505937099 * s_c
+    b_ok = 0.0259040371 * l_c + 0.7827717662 * m_c - 0.8086757660 * s_c
+
+    return (L, a, b_ok)
+
+
+def oklab_to_srgb(L: float, a: float, b_ok: float) -> Tuple[float, float, float]:
+    """Convert OKLab to sRGB [0,1]."""
+    l_c = L + 0.3963377774 * a + 0.2158037573 * b_ok
+    m_c = L - 0.1055613458 * a - 0.0638541728 * b_ok
+    s_c = L - 0.0894841775 * a - 1.2914855480 * b_ok
+
+    l_ = l_c * l_c * l_c
+    m_ = m_c * m_c * m_c
+    s_ = s_c * s_c * s_c
+
+    r_lin = +4.0767416621 * l_ - 3.3077115913 * m_ + 0.2309699292 * s_
+    g_lin = -1.2684380046 * l_ + 2.6097574011 * m_ - 0.3413193965 * s_
+    b_lin = -0.0041960863 * l_ - 0.7034186147 * m_ + 1.7076147010 * s_
+
+    r = clamp(linear_to_srgb(clamp(r_lin, 0, 1)), 0, 1)
+    g = clamp(linear_to_srgb(clamp(g_lin, 0, 1)), 0, 1)
+    b = clamp(linear_to_srgb(clamp(b_lin, 0, 1)), 0, 1)
+
+    return (r, g, b)
+
+
+# ── HSL ↔ RGB ──
+
+def hsl_to_rgb(h_deg: float, s_pct: float, l_pct: float) -> Tuple[float, float, float]:
+    """Convert HSL (degrees, percent, percent) to RGB [0,1]."""
+    h = h_deg / 360.0
+    s = s_pct / 100.0
+    l = l_pct / 100.0
+
+    if s == 0:
+        return (l, l, l)
+
+    def hue_to_rgb(p, q, t):
+        if t < 0: t += 1
+        if t > 1: t -= 1
+        if t < 1/6: return p + (q - p) * 6 * t
+        if t < 1/2: return q
+        if t < 2/3: return p + (q - p) * (2/3 - t) * 6
+        return p
+
+    q = l * (1 + s) if l < 0.5 else l + s - l * s
+    p = 2 * l - q
+
+    r = hue_to_rgb(p, q, h + 1/3)
+    g = hue_to_rgb(p, q, h)
+    b = hue_to_rgb(p, q, h - 1/3)
+
+    return (r, g, b)
+
+
+def rgb_to_hsl(r: float, g: float, b: float) -> Tuple[float, float, float]:
+    """Convert RGB [0,1] to HSL (degrees, percent, percent)."""
+    max_c = max(r, g, b)
+    min_c = min(r, g, b)
+    l = (max_c + min_c) / 2.0
+
+    if max_c == min_c:
+        h = s = 0.0
+    else:
+        d = max_c - min_c
+        s = d / (2.0 - max_c - min_c) if l > 0.5 else d / (max_c + min_c)
+
+        if max_c == r:
+            h = (g - b) / d + (6 if g < b else 0)
+        elif max_c == g:
+            h = (b - r) / d + 2
+        else:
+            h = (r - g) / d + 4
+
+        h /= 6.0
+
+    return (h * 360.0, s * 100.0, l * 100.0)
+
+
+# ── OKLab Operations ──
+
+def rotate_ab(a: float, b: float, degrees: float) -> Tuple[float, float]:
+    """Rotate hue in OKLab a-b plane by given degrees."""
+    rad = math.radians(degrees)
+    cos_r = math.cos(rad)
+    sin_r = math.sin(rad)
+    return (a * cos_r - b * sin_r, a * sin_r + b * cos_r)
+
+
+def set_chroma(a: float, b: float, target_c: float) -> Tuple[float, float]:
+    """Set the chroma (magnitude in a-b plane) to target value."""
+    current_c = math.sqrt(a * a + b * b)
+    if current_c < 1e-10:
+        return (target_c, 0.0)  # Default direction
+    scale = target_c / current_c
+    return (a * scale, b * scale)
+
+
+def get_chroma(a: float, b: float) -> float:
+    """Get chroma magnitude from a-b values."""
+    return math.sqrt(a * a + b * b)
+
+
+def compute_delta_e_oklab(
+    L1: float, a1: float, b1: float,
+    L2: float, a2: float, b2: float,
+) -> float:
+    """Compute ΔE in OKLab space (perceptual color difference)."""
+    return math.sqrt((L1 - L2) ** 2 + (a1 - a2) ** 2 + (b1 - b2) ** 2)
+
+
+# ── Batch Operations (PyTorch) ──
+
+def hsl_to_oklab_batch(hsl: torch.Tensor) -> torch.Tensor:
+    """
+    Batch convert HSL [0,1] normalized to OKLab.
+
+    Args:
+        hsl: (..., 3) tensor with H,S,L in [0,1]
+
+    Returns:
+        (..., 3) tensor with L,a,b in OKLab
+    """
+    h = hsl[..., 0] * 360.0  # Back to degrees
+    s = hsl[..., 1] * 100.0  # Back to percent
+    l = hsl[..., 2] * 100.0  # Back to percent
+
+    # HSL to RGB (vectorized)
+    h_norm = h / 360.0
+    q = torch.where(l / 100.0 < 0.5,
+                    (l / 100.0) * (1 + s / 100.0),
+                    (l / 100.0) + (s / 100.0) - (l / 100.0) * (s / 100.0))
+    p = 2 * (l / 100.0) - q
+
+    def hue2rgb(p, q, t):
+        t = t % 1.0
+        r = torch.where(t < 1/6, p + (q - p) * 6 * t,
+            torch.where(t < 1/2, q,
+            torch.where(t < 2/3, p + (q - p) * (2/3 - t) * 6, p)))
+        return r
+
+    r = hue2rgb(p, q, h_norm + 1/3)
+    g = hue2rgb(p, q, h_norm)
+    b = hue2rgb(p, q, h_norm - 1/3)
+
+    # Handle achromatic (s == 0)
+    achromatic = (s < 0.001)
+    r = torch.where(achromatic, l / 100.0, r)
+    g = torch.where(achromatic, l / 100.0, g)
+    b = torch.where(achromatic, l / 100.0, b)
+
+    # sRGB to linear
+    r_lin = torch.where(r <= 0.04045, r / 12.92, ((r + 0.055) / 1.055) ** 2.4)
+    g_lin = torch.where(g <= 0.04045, g / 12.92, ((g + 0.055) / 1.055) ** 2.4)
+    b_lin = torch.where(b <= 0.04045, b / 12.92, ((b + 0.055) / 1.055) ** 2.4)
+
+    # Linear RGB to OKLab
+    l_ = 0.4122214708 * r_lin + 0.5363325363 * g_lin + 0.0514459929 * b_lin
+    m_ = 0.2119034982 * r_lin + 0.6806995451 * g_lin + 0.1073969566 * b_lin
+    s_ = 0.0883024619 * r_lin + 0.2817188376 * g_lin + 0.6299787005 * b_lin
+
+    l_c = torch.sign(l_) * torch.abs(l_).pow(1/3)
+    m_c = torch.sign(m_) * torch.abs(m_).pow(1/3)
+    s_c = torch.sign(s_) * torch.abs(s_).pow(1/3)
+
+    L_ok = 0.2104542553 * l_c + 0.7936177850 * m_c - 0.0040720468 * s_c
+    a_ok = 1.9779984951 * l_c - 2.4285922050 * m_c + 0.4505937099 * s_c
+    b_ok = 0.0259040371 * l_c + 0.7827717662 * m_c - 0.8086757660 * s_c
+
+    return torch.stack([L_ok, a_ok, b_ok], dim=-1)
+
+
+def denormalize_hsl(hsl_norm: torch.Tensor) -> torch.Tensor:
+    """Convert normalized HSL [0,1] to degrees/percent format."""
+    result = hsl_norm.clone()
+    result[..., 0] *= 360.0  # H: [0,1] → [0,360]
+    result[..., 1] *= 100.0  # S: [0,1] → [0,100]
+    result[..., 2] *= 100.0  # L: [0,1] → [0,100]
+    return result
+
+
+class OKLabMSELoss(nn.Module):
+    """
+    Perceptually uniform loss in OKLab space.
+
+    Converts predicted and target HSL values to OKLab, then computes MSE.
+    This handles hue circularity correctly (359° ≈ 1°) because OKLab
+    represents hue as a-b coordinates, not an angle.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+        self,
+        pred_hsl: torch.Tensor,    # (B, 3) predicted HSL in [0,1]
+        target_hsl: torch.Tensor,  # (B, 3) target HSL in [0,1]
+    ) -> torch.Tensor:
+        """Compute perceptually uniform loss."""
+        pred_oklab = hsl_to_oklab_batch(pred_hsl)
+        target_oklab = hsl_to_oklab_batch(target_hsl)
+
+        return torch.nn.functional.mse_loss(pred_oklab, target_oklab)