From a4dce23ef8c0a936c15a9bdddf8b1c6f1d50a625 Mon Sep 17 00:00:00 2001 From: Andrii Ryzhkov Date: Sat, 16 May 2026 16:54:40 +0200 Subject: [PATCH] Fix wrong colors on non-RGGB sensors via RGGB origin shift and daylight WB round-trip --- examples/ai_denoise.lua | 4 +- examples/ai_raw_denoise.lua | 107 +++++++++++++++++++++++++++++------- 2 files changed, 90 insertions(+), 21 deletions(-) diff --git a/examples/ai_denoise.lua b/examples/ai_denoise.lua index 5d66b079..1528d6f3 100644 --- a/examples/ai_denoise.lua +++ b/examples/ai_denoise.lua @@ -312,9 +312,9 @@ dt.register_lib( local script_data = {} script_data.metadata = { - name = "ai_denoise", + name = "AI denoise", purpose = "tile-based AI denoise using the darktable.ai Lua API", - author = "Andrii Ryzhkov ", + author = "Andrii Ryzhkov", help = "https://docs.darktable.org/lua/stable/lua.scripts.manual/scripts/examples/ai_denoise" } diff --git a/examples/ai_raw_denoise.lua b/examples/ai_raw_denoise.lua index ec5ea46e..b980f678 100644 --- a/examples/ai_raw_denoise.lua +++ b/examples/ai_raw_denoise.lua @@ -91,6 +91,37 @@ local OVERLAP = 32 -- F U N C T I O N S -- - - - - - - - - - - - - - - - - - - - - - - - +-- Bayer-only CFA channel lookup. Returns 0/1/2 (R/G/B) for the cell +-- at sensor position (r, c) given the 32-bit `filters` mask from +-- load_raw's meta table. Used to find the R-cell origin so bayer_pack's +-- plane 0 always holds R. X-Trans (filters == 9) uses meta.xtrans +-- instead and is not handled here. +local function fc(r, c, filters) + local shift = ((r % 2) * 4) + ((c % 2) * 2) + return math.floor(filters / (2 ^ shift)) % 4 +end + +-- daylight WB multipliers (R, G=1, B) derived from the camera's +-- XYZ->camRGB color matrix and the D65 white point. mirrors the C +-- path's _bayer_wb_daylight in restore_raw_bayer.c — the rawdenoise +-- bayer model is trained on daylight-WB'd input, so feeding it +-- un-WB'd camRGB produces wrong colors. Falls back to {1, 1, 1} if +-- the matrix is missing or degenerate. +local function daylight_wb(color_matrix) + local d65 = {0.95047, 1.0, 1.08883} + if not color_matrix then return 1.0, 1.0, 1.0 end + local resp = {0, 0, 0} + for c = 1, 3 do + resp[c] = color_matrix[c][1] * d65[1] + + color_matrix[c][2] * d65[2] + + color_matrix[c][3] * d65[3] + end + if resp[1] <= 0 or resp[2] <= 0 or resp[3] <= 0 then + return 1.0, 1.0, 1.0 + end + return resp[2] / resp[1], 1.0, resp[2] / resp[3] +end + -- Surround the input with `pad` pixels on every side, replicating the -- outermost rows/columns ("edge clamp" padding). Boundary tiles then -- see plausible context instead of a literal image edge, so the @@ -221,22 +252,47 @@ local function denoise_one(ctx, img, tile_size, job, base_pct, span_pct) local cfa, meta = dt.ai.load_raw(img) if not cfa then return false, "load_raw failed" end - -- crop off optical-black / masked pixels (Canon and others ship - -- the full sensor buffer including non-light-sensing strips on - -- the top/left). visible_* and crop_* come from the image - -- metadata; we snap crop offsets down to an even pixel so the - -- 2x2 CFA phase at the visible-region origin matches the buffer - -- (otherwise bayer_pack would group the wrong colours) - if meta.crop_x and meta.crop_y - and (meta.crop_x > 0 or meta.crop_y > 0) + -- this script is Bayer-only: filters == 9 is X-Trans (handled by a + -- different preprocessing pipeline + model_linear.onnx); filters == 0 + -- means monochrome / non-CFA, which bayer_pack can't handle either + if not meta.filters or meta.filters == 0 or meta.filters == 9 then + return false, "unsupported CFA pattern (Bayer-only script)" + end + + -- Crop to the visible region and shift the origin onto the R cell + -- of the CFA pattern. bayer_pack splits by position only, so to make + -- plane 0 always hold R (which the model expects) the cropped CFA's + -- (0, 0) must sit at an R site. Mirrors the C path's _bayer_origin + -- + RGGB shift in restore_raw_bayer.c. + local cy = 0 + local cx = 0 + local cw = cfa:shape()[4] + local ch = cfa:shape()[3] + if meta.crop_x and meta.crop_y then + cy = meta.crop_y - (meta.crop_y % 2) + cx = meta.crop_x - (meta.crop_x % 2) + cw = (meta.visible_width or cw) + (meta.crop_x - cx) + ch = (meta.visible_height or ch) + (meta.crop_y - cy) + end + -- find R cell within the 2x2 starting at (cy, cx) + local sy, sx = 0, 0 + for ty = 0, 1 do + for tx = 0, 1 do + if fc(cy + ty, cx + tx, meta.filters) == 0 then + sy, sx = ty, tx + end + end + end + cy = cy + sy + cx = cx + sx + ch = ch - sy + cw = cw - sx + -- ensure even dims for bayer_pack + ch = ch - (ch % 2) + cw = cw - (cw % 2) + if cy > 0 or cx > 0 + or cw < cfa:shape()[4] or ch < cfa:shape()[3] then - local cx = meta.crop_x - (meta.crop_x % 2) - local cy = meta.crop_y - (meta.crop_y % 2) - local cw = meta.visible_width + (meta.crop_x - cx) - local ch = meta.visible_height + (meta.crop_y - cy) - -- ensure even dims for bayer_pack - cw = cw - (cw % 2) - ch = ch - (ch % 2) cfa = cfa:crop(cy, cx, ch, cw) end @@ -244,15 +300,23 @@ local function denoise_one(ctx, img, tile_size, job, base_pct, span_pct) -- levels across the 4 CFA sites, in which case a single scalar -- (val - black) / (white - black) normalisation is exact. cameras -- with non-uniform per-site black levels (e.g. some PDAF sensors) - -- would need per-plane scaling -- left as an exercise + -- would need per-plane scaling via scale_add_planes -- left as an + -- exercise local black = meta.black_level local range = meta.white_level - black if range <= 0 then range = 65535 end -- pack the 1-channel CFA into 4 phase planes so each channel holds - -- one Bayer site type; the model expects [1,4,H/2,W/2] input + -- one Bayer site type. After the RGGB shift above: + -- plane 0 = R, planes 1,2 = G, plane 3 = B local packed = cfa:bayer_pack() packed:scale_add(1.0 / range, -black / range) + + -- apply daylight WB so the four planes sit in the same range the + -- model was trained on (R, G, G, B → R*wb_R, G, G, B*wb_B). without + -- this, the model sees unbalanced channels and produces wrong colors + local wb_R, wb_G, wb_B = daylight_wb(meta.color_matrix) + packed:scale_add_planes({wb_R, wb_G, wb_G, wb_B}) local pH = packed:shape()[3] local pW = packed:shape()[4] @@ -286,6 +350,11 @@ local function denoise_one(ctx, img, tile_size, job, base_pct, span_pct) out_rgb:scale_add(in_mean / out_mean) end + -- invert the daylight WB so the saved tensor is un-WB'd camRGB + -- (save_dng_linear's contract — the consumer applies AsShotNeutral + -- on import) + out_rgb:scale_add_planes({1.0 / wb_R, 1.0 / wb_G, 1.0 / wb_B}) + local base = img.filename:match("(.+)%..+$") or img.filename local out_path = img.path .. "/" .. base .. "_rawdenoised.dng" -- LinearRaw DNG: tensor values must be normalised to [0, 1] camRGB; @@ -401,9 +470,9 @@ dt.register_lib( local script_data = {} script_data.metadata = { - name = "ai_raw_denoise", + name = "AI raw denoise", purpose = "tile-based AI raw denoise using the darktable.ai Lua API", - author = "Andrii Ryzhkov ", + author = "Andrii Ryzhkov", help = "https://docs.darktable.org/lua/stable/lua.scripts.manual/scripts/examples/ai_raw_denoise" }