DarkRoom/core/dr-gpu/tests/hot_pixels.rs

//! TRACES: FR-RAW-3
//! Hot and dead photosite repair, end to end on a device.
//!
//! Each test renders a frame twice — once with a defect, once without — and
//! compares the finished pixels. That is the only comparison that means
//! anything: the repair happens on the mosaic, and what a photographer would
//! see of a defect it missed is the coloured cross the demosaic makes of it.

use dr_decode::{CfaPattern, CropRect, RawImage};
use dr_gpu::{AdjustPass, Demosaicer, GpuContext};
use dr_pipeline::EditGraph;

const SIZE: u32 = 36;
const WHITE: u16 = 4095;

fn ctx() -> Option<GpuContext> {
    pollster::block_on(GpuContext::new_headless()).ok()
}

/// A flat frame at `level`, with `set` applied to its photosites.
fn frame(pattern: CfaPattern, level: u16, set: &[(u32, u32, u16)]) -> RawImage {
    let mut data = vec![level; (SIZE * SIZE) as usize];
    for &(x, y, v) in set {
        data[(y * SIZE + x) as usize] = v;
    }
    RawImage {
        width: SIZE,
        height: SIZE,
        data,
        cfa_pattern: pattern,
        black_level: [0; 4],
        white_level: WHITE,
        wb_coeffs: [1.0, 1.0, 1.0, 1.0],
        color_matrix: Some([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]),
        samples_per_pixel: 1,
        profile: None,
        profile_tables: None,
        make: String::new(),
        model: String::new(),
        crop: CropRect {
            x: 0,
            y: 0,
            width: SIZE,
            height: SIZE,
        },
    }
}

fn render(ctx: &GpuContext, raw: &RawImage) -> Vec<u8> {
    let source = Demosaicer::new(ctx)
        .expect("demosaicer")
        .run(raw)
        .expect("demosaic");
    let shader = EditGraph::default_chain().compose();
    let mut adjust = AdjustPass::new(ctx);
    adjust.render(&source, &shader, SIZE, SIZE).expect("render");
    adjust.export_pixels().expect("readback").0
}

/// The largest channel difference between two renders.
fn worst(a: &[u8], b: &[u8]) -> u8 {
    a.iter().zip(b).map(|(x, y)| x.abs_diff(*y)).max().unwrap()
}

const MIDDLE: u32 = SIZE / 2;

/// **The feature.** A photosite at white in a dark frame — a hot pixel in a
/// night sky — leaves no trace in the rendered picture.
#[test]
fn a_hot_photosite_in_a_dark_frame_is_invisible() {
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let clean = render(&ctx, &frame(CfaPattern::Rggb, 40, &[]));
    for (x, y) in [
        (MIDDLE, MIDDLE),
        (MIDDLE + 1, MIDDLE),
        (MIDDLE + 1, MIDDLE + 1),
    ] {
        let hot = render(&ctx, &frame(CfaPattern::Rggb, 40, &[(x, y, WHITE)]));
        let diff = worst(&clean, &hot);
        assert!(
            diff <= 1,
            "a hot photosite at ({x}, {y}) still shows, by {diff}"
        );
    }
}

/// The same for one stuck dark in a lit area.
#[test]
fn a_dead_photosite_in_a_lit_frame_is_invisible() {
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let clean = render(&ctx, &frame(CfaPattern::Rggb, 1600, &[]));
    let dead = render(&ctx, &frame(CfaPattern::Rggb, 1600, &[(MIDDLE, MIDDLE, 0)]));
    let diff = worst(&clean, &dead);
    assert!(diff <= 1, "a dead photosite still shows, by {diff}");
}

/// **What it must not eat.** A point of real light lands on a patch of
/// photosites, not one — so a 3×3 highlight survives, even at its brightest.
#[test]
fn a_small_real_highlight_survives() {
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let mut star = Vec::new();
    for dy in 0..3 {
        for dx in 0..3 {
            star.push((MIDDLE - 1 + dx, MIDDLE - 1 + dy, WHITE));
        }
    }
    let clean = render(&ctx, &frame(CfaPattern::Rggb, 40, &[]));
    let lit = render(&ctx, &frame(CfaPattern::Rggb, 40, &star));
    let at = ((MIDDLE * SIZE + MIDDLE) * 4 + 1) as usize;
    assert!(
        lit[at] > clean[at] + 100,
        "the highlight was repaired away: {} against a background of {}",
        lit[at],
        clean[at]
    );
}

/// The Fujifilm path goes through the same repair, with its own tile.
#[test]
fn a_hot_photosite_on_x_trans_is_invisible() {
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let clean = render(&ctx, &frame(CfaPattern::XTrans, 40, &[]));
    let hot = render(
        &ctx,
        &frame(CfaPattern::XTrans, 40, &[(MIDDLE, MIDDLE, WHITE)]),
    );
    let diff = worst(&clean, &hot);
    assert!(diff <= 1, "a hot X-Trans photosite still shows, by {diff}");
}

/// The repair alone, read back (FR-DEV-3g): the learned demosaic takes the
/// mosaic this pass leaves, so it must be the same pass and nothing more —
/// the hot photosite replaced, a real highlight and every other photosite
/// untouched.
#[test]
fn the_repaired_mosaic_reads_back_with_only_the_defect_changed() {
    let Some(ctx) = ctx() else {
        eprintln!("no GPU adapter; skipping");
        return;
    };
    let d = Demosaicer::new(&ctx).expect("demosaicer");
    let mut star = vec![(MIDDLE, MIDDLE, WHITE)];
    for dy in 0..3 {
        for dx in 0..3 {
            star.push((4 + dx, 4 + dy, WHITE));
        }
    }
    let before = frame(CfaPattern::Rggb, 40, &star);
    let mut raw = before.clone();
    let changed = d.repair_hot_pixels(&mut raw).expect("repair");
    assert_eq!(changed, 1, "only the lone hot photosite should change");
    let at = (MIDDLE * SIZE + MIDDLE) as usize;
    assert_eq!(
        raw.data[at], 40,
        "repaired to its brightest same-colour neighbour"
    );
    let others = (0..raw.data.len()).filter(|&i| i != at);
    assert!(others.into_iter().all(|i| raw.data[i] == before.data[i]));
}