DarkRoom/ui/dr-ui/src/develop/mod.rs

//! The develop session — capabilities in, rendered image out.
//!
//! This is the only place the UI touches the pipeline, and it does so through
//! two calls: [`dr_pipeline::EditGraph::capabilities`] to learn what controls
//! to build, and `set_param` to change one. It never names an operation, and
//! it knows nothing about shaders.
//!
//! Whether a control is a slider or a switch follows from the parameter's
//! declared [`ParamKind`], not from which parameter it is (ARCH §4.3), so a
//! new operation appears in the panel with no change here (FR-DEV-3c).
//!
//! Split into one module per area of the session (framing, masks,
//! segmentation, rendering, ...) rather than kept as one file — see
//! `docs/dev/code-health.md` CH-1. `DevelopSession`'s fields are `pub(super)`
//! so its `impl` blocks can live beside the area of behaviour they belong to
//! instead of all in one place; nothing outside this module sees them, since
//! only methods were ever exported.

mod curves;
mod denoise;
mod denoise_cache;
mod framing;
mod history;
mod mask_ops;
mod masks;
mod render;
mod repairs;
mod rows;
mod segmentation;
mod session;
mod tabs;
mod white_balance;

pub use denoise::DenoiseStatus;
pub use framing::CropAspect;
pub use masks::MASK_COLOURS;
pub use segmentation::{Abandon, RefinedInstance, Segmented, SessionId};
pub use session::DevelopSession;

/// Longest edge the model and the masks work at.
///
/// ~1.3 MP at 3:2. Large enough that an outline is within a pixel or two of
/// where it belongs, small enough that a distance transform over it is a few
/// milliseconds and its field a few megabytes.
pub(super) const SEGMENT_PROXY_EDGE: u32 = 1600;

/// Test-only helpers shared by more than one of this module's submodules.
///
/// `headless`, `read_back` and `grey_session` were each defined once in the
/// pre-split file and called from tests all over it. Splitting the tests with
/// the code they exercise left these three needed in most of the resulting
/// files, so they live here once instead of being copied.
#[cfg(test)]
pub(crate) mod test_support {
    use super::*;
    use dr_gpu::GpuContext;

    /// TRACES: FR-DSP-1 | AC-8
    /// Copy a displayed frame back to the CPU, for assertions and nothing else.
    ///
    /// The library has no such function on purpose: S1 removed the display
    /// readback, and AC-8 is the assertion that it stayed removed. A test that
    /// wants to look at the pixels therefore has to do the copy itself, which
    /// is exactly the right shape — the round-trip lives in the test binary
    /// and cannot be reached from a shipping one.
    ///
    /// Doubles as the proof: this only compiles because the image *is* a wgpu
    /// texture. Hand it a `SharedPixelBuffer`-backed image and it panics.
    pub(super) fn read_back(ctx: &GpuContext, image: &slint::Image) -> Vec<u8> {
        let texture = image
            .to_wgpu_29_texture()
            .expect("the develop canvas must be a GPU texture, not a pixel buffer");
        let (w, h) = (texture.width(), texture.height());

        // Buffer rows must be aligned to COPY_BYTES_PER_ROW_ALIGNMENT.
        let unpadded = w * 4;
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let padded = unpadded.div_ceil(align) * align;

        let buf = ctx.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("test-readback"),
            size: u64::from(padded * h),
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });

        let mut enc = ctx.device.create_command_encoder(&Default::default());
        enc.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: &texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &buf,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(padded),
                    rows_per_image: Some(h),
                },
            },
            wgpu::Extent3d {
                width: w,
                height: h,
                depth_or_array_layers: 1,
            },
        );
        ctx.queue.submit(Some(enc.finish()));

        let slice = buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |r| {
            let _ = tx.send(r);
        });
        ctx.device
            .poll(wgpu::PollType::wait_indefinitely())
            .expect("poll");
        rx.recv().expect("map").expect("map");

        let data = slice.get_mapped_range();
        let mut out = Vec::with_capacity((unpadded * h) as usize);
        for row in 0..h {
            let start = (row * padded) as usize;
            out.extend_from_slice(&data[start..start + unpadded as usize]);
        }
        drop(data);
        buf.unmap();
        out
    }

    /// One device for the whole test binary.
    ///
    /// This opened a *new* `GpuContext` per test, and `cargo test` runs tests
    /// on as many threads as there are cores — so a full run asked the driver
    /// to bring up a dozen Vulkan devices at once and the binary died with
    /// SIGSEGV. Serially it passed, which is what made it look like flakiness
    /// rather than a bug in the harness.
    ///
    /// A `GpuContext` is an `Arc<Device>` and an `Arc<Queue>`, so sharing one
    /// is a refcount rather than a copy, and wgpu is explicit that both are
    /// safe to use from several threads. Nothing here mutates the context; the
    /// per-test state is in the passes and the sessions built on top of it.
    ///
    /// `OnceLock` rather than `lazy_static`: the initialiser runs once however
    /// many threads arrive together, and the losers block until it is done —
    /// which is precisely the property that was missing.
    pub(crate) fn headless() -> Option<GpuContext> {
        static SHARED: std::sync::OnceLock<Option<GpuContext>> = std::sync::OnceLock::new();
        SHARED
            .get_or_init(|| pollster::block_on(dr_gpu::GpuContext::new_headless()).ok())
            .clone()
    }

    /// A flat grey session with nothing segmented, and its render.
    pub(super) fn grey_session(ctx: &GpuContext) -> (DevelopSession, Vec<u8>) {
        let rgba: Vec<u8> = (0..64 * 64).flat_map(|_| [128u8, 128, 128, 255]).collect();
        let mut session =
            DevelopSession::open_rgb(ctx, &rgba, 64, 64, dr_types::Orientation::NORMAL)
                .expect("session");
        assert!(
            !session.has_segmentation(),
            "the premise: no model has been run"
        );
        let before = read_back(ctx, &session.render(64, 64).expect("render"));
        (session, before)
    }

    /// A session holding a segmentation with one instance over the left half.
    ///
    /// Built rather than detected. What these tests need is coverage of a
    /// *known* shape, so that "the stored raster renders what the model's did"
    /// is a comparison rather than a hope — and asking a real run what it
    /// happened to find in a synthetic frame would make the assertion depend
    /// on the weights.
    pub(super) fn session_with_a_left_half_subject(ctx: &GpuContext) -> DevelopSession {
        let rgba: Vec<u8> = (0..64 * 64).flat_map(|_| [128u8, 128, 128, 255]).collect();
        let mut session =
            DevelopSession::open_rgb(ctx, &rgba, 64, 64, dr_types::Orientation::NORMAL)
                .expect("session");

        let (pw, ph) = session.mask_raster_size();
        let (pw, ph) = (pw as usize, ph as usize);
        let mut mask = vec![0u8; pw * ph];
        for y in 0..ph {
            for x in 0..pw / 2 {
                mask[y * pw + x] = 255;
            }
        }

        session.segmentation = Some(crate::segmentation::Segmentation::for_test(
            vec![crate::segmentation::InstanceSummary {
                class_name: "dog".into(),
                score: 0.9,
                mask,
                bbox: (0.0, 0.0, (pw / 2) as f32, ph as f32),
            }],
            Vec::new(),
            0xfeed,
            (pw, ph),
        ));
        session.subjects = None;
        session.subject_key = 0;
        session
    }
}