new color rendering

app.py

@@ -147,8 +147,8 @@ def _discard(path):
         pass
 
 
-# Decoded image (+ month) per upload, RAM only. Lets the user re-run with a
-# different depth without re-reading the file (which we discard after decoding).
+# Raw object (+ month) per upload, RAM only. Lets the user re-run with a different
+# depth without re-reading the file (which we discard right after opening).
 _DECODE_CACHE = {}
 _DECODE_CACHE_MAX = 6
 
@@ -156,7 +156,11 @@ _DECODE_CACHE_MAX = 6
 def _cache_put(path, value):
     _DECODE_CACHE[path] = value
     while len(_DECODE_CACHE) > _DECODE_CACHE_MAX:
-        _DECODE_CACHE.pop(next(iter(_DECODE_CACHE)))
+        old = _DECODE_CACHE.pop(next(iter(_DECODE_CACHE)))
+        try:
+            old[0].close()  # free the evicted RawPy's buffers
+        except Exception:
+            pass
 
 
 def on_upload(path):
@@ -172,22 +176,22 @@ def correct(path, lat, lng, depth):
     if lat is None or lng is None:
         raise gr.Error("Set a location on the map.")
 
-    # Decode once per upload, cache it, then discard the file. Re-runs (e.g. a
-    # new depth) reuse the cached image — no second read of a deleted file.
+    # imread once per upload, cache the raw object, then discard the file. Re-runs
+    # (e.g. a new depth) re-render from the in-memory raw — no second file read.
     cached = _DECODE_CACHE.get(path)
     if cached is None:
         if not os.path.exists(path):
             raise gr.Error("Please re-upload the RAW — the previous file was cleared.")
         month = exif_gps.extract_month(path)  # read metadata before discarding
         try:
-            linear = color.decode_raw(path)
+            raw = color.open_raw(path)
         except Exception as exc:
             raise gr.Error(f"Could not decode RAW ({type(exc).__name__}): {exc}")
         finally:
             _discard(path)
-        cached = (linear, month)
+        cached = (raw, month)
         _cache_put(path, cached)
-    linear, month = cached  # month is None when no EXIF date -> engine uses yearly data
+    raw, month = cached  # month is None when no EXIF date -> engine uses yearly data
 
     try:
         illum = client.fetch_illuminant(lat, lng, depth, month)
@@ -195,9 +199,8 @@ def correct(path, lat, lng, depth):
         raise gr.Error(f"Parameter API error: {exc}")
 
     print(f"[uwcolorkit] month={month} stub={client.using_stub()}")
-    matrix = color.correction_matrix(illum.as_tuple())
-    before = color.encode_display(linear)
-    after = color.encode_display(color.apply_correction(linear, matrix))
+    before = color.render(raw, color.baseline_wb(raw))
+    after = color.render(raw, color.wb_from_chromaticity(raw, illum.as_tuple()))
     return (before, after)
 
 

color.py

@@ -1,47 +1,72 @@
-"""RAW decode and Bradford chromatic adaptation."""
+"""RAW decode + camera-space white balance from an illuminant chromaticity.
 
-import numpy as np
-
-D65_XY = (0.3127, 0.3290)
+White balance is applied the way the iOS app does it (`CIRAWFilter.neutralChromaticity`):
+as per-channel multipliers in the camera's *native* space, derived from the target
+illuminant through the camera's own XYZ->camera matrix. The earlier Bradford-CAT-in-sRGB
+approach used the wrong operator in the wrong space and produced a desaturated,
+magenta/hazy result; this neutralises the illuminant the way the app does.
 
-M_RGB2XYZ = np.array([
-    [0.4124564, 0.3575761, 0.1804375],
-    [0.2126729, 0.7151522, 0.0721750],
-    [0.0193339, 0.1191920, 0.9503041],
-])
-M_XYZ2RGB = np.linalg.inv(M_RGB2XYZ)
+This is deliberately a *white-balance* preview, not a full ISP: no auto-exposure, tone
+curve, or contrast. 'before' and 'after' share identical decode settings and differ
+only in the WB multipliers, so the comparison shows the correction alone.
+"""
 
-M_BRADFORD = np.array([
-    [0.8951000, 0.2664000, -0.1614000],
-    [-0.7502000, 1.7135000, 0.0367000],
-    [0.0389000, -0.0685000, 1.0296000],
-])
-M_BRADFORD_INV = np.linalg.inv(M_BRADFORD)
+import numpy as np
 
 HALF_SIZE = True
 MAX_LONG_EDGE = 2000
 
+# Shared postprocess: linear 16-bit, no auto-exposure, half-res. WB + output colour
+# are supplied per call so before/after differ only by white balance.
+_POST = dict(output_bps=16, gamma=(1, 1), no_auto_bright=True, half_size=HALF_SIZE)
+
 
 def xy_to_XYZ(x, y, Y=1.0):
     return np.array([Y * x / y, Y, Y * (1.0 - x - y) / y])
 
 
-def bradford_cat(src_white_XYZ, dst_white_XYZ):
-    src = M_BRADFORD @ src_white_XYZ
-    dst = M_BRADFORD @ dst_white_XYZ
-    return M_BRADFORD_INV @ np.diag(dst / src) @ M_BRADFORD
+def open_raw(path):
+    """imread and return the RawPy object. The caller may delete the file right
+    after — the unpacked sensor data stays in RAM, so depth re-runs re-render from
+    memory (no second read) and the upload is still discarded for privacy."""
+    import rawpy
+
+    return rawpy.imread(str(path))
+
 
+def wb_from_chromaticity(raw, illuminant_xy):
+    """Raw WB multipliers [R, G, B, G2] that render `illuminant_xy` neutral, derived
+    from the camera's XYZ->camera matrix (`raw.rgb_xyz_matrix`). This is the camera-
+    space equivalent of CIRAWFilter.neutralChromaticity. Validated: feeding D65 here
+    reproduces `raw.daylight_whitebalance` to 4 decimals."""
+    M = np.asarray(raw.rgb_xyz_matrix, dtype=float)[:3]   # XYZ -> camera
+    cam = M @ xy_to_XYZ(*illuminant_xy)
+    cam = np.where(np.abs(cam) < 1e-9, 1e-9, cam)
+    mul = 1.0 / cam
+    mul = mul / mul[1]                                    # normalise green = 1
+    return [float(mul[0]), 1.0, float(mul[2]), 1.0]
 
-def correction_matrix(illuminant_xy, target_xy=D65_XY):
-    cat = bradford_cat(xy_to_XYZ(*illuminant_xy), xy_to_XYZ(*target_xy))
-    return M_XYZ2RGB @ cat @ M_RGB2XYZ
 
+def baseline_wb(raw):
+    """'before' WB — the camera's as-shot balance (the photo as captured), falling
+    back to the daylight balance. Green-normalised."""
+    cw = np.asarray(raw.camera_whitebalance, dtype=float)
+    if cw.size < 3 or cw[1] <= 0 or not np.all(np.isfinite(cw[:3])):
+        cw = np.asarray(raw.daylight_whitebalance, dtype=float)
+    cw = cw / cw[1]
+    return [float(cw[0]), 1.0, float(cw[2]), 1.0]
 
-def apply_correction(linear_rgb, matrix):
-    return linear_rgb @ matrix.T
 
+def render(raw, wb, max_long_edge=MAX_LONG_EDGE):
+    """postprocess `raw` with WB multipliers `wb` -> uint8 sRGB display image."""
+    import rawpy
+
+    rgb16 = raw.postprocess(user_wb=list(wb), output_color=rawpy.ColorSpace.sRGB, **_POST)
+    linear = _resize_max(rgb16.astype(np.float32) / 65535.0, max_long_edge)
+    return _encode_display(linear)
 
-def encode_display(linear_rgb):
+
+def _encode_display(linear_rgb):
     lin = np.clip(linear_rgb, 0.0, 1.0)
     srgb = np.where(lin <= 0.0031308, lin * 12.92, 1.055 * np.power(lin, 1 / 2.4) - 0.055)
     return (np.clip(srgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
@@ -60,20 +85,3 @@ def _resize_max(img, max_long_edge):
         for c in range(img.shape[2])
     ]
     return np.stack(chans, axis=-1)
-
-
-def decode_raw(path, half_size=HALF_SIZE, max_long_edge=MAX_LONG_EDGE):
-    import rawpy
-
-    # rawpy.imread (Cython) needs an exact str; Gradio passes a NamedString.
-    with rawpy.imread(str(path)) as raw:
-        rgb16 = raw.postprocess(
-            no_auto_bright=True,
-            output_bps=16,
-            use_camera_wb=False,
-            use_auto_wb=False,
-            gamma=(1, 1),
-            output_color=rawpy.ColorSpace.sRGB,
-            half_size=half_size,
-        )
-    return _resize_max(rgb16.astype(np.float32) / 65535.0, max_long_edge)