DarkRoom/core/dr-pano/src/xfeat.rs

//! TRACES: FR-MRG-8
//! The XFeat detector — the network under tract, and the decoder after it.
//!
//! Apache-2.0 weights (`models/LICENCE.md`), exported at a fixed shape by
//! `tools/export-xfeat.sh` and loaded through the same `ort`-over-tract
//! backend `dr-segment` and `dr-face` use, so this adds no runtime and no C
//! to the tree. ~300 ms per frame on the reference desktop, ~400 ms on the
//! tablet (S15.2, S15.4).

use crate::features::{decode_xfeat, DecodeOptions, Features, XFeatMaps, DESCRIPTOR_LEN};
use crate::image::Gray;
use crate::PanoError;

/// The two input shapes the shipped exports were made for: one landscape,
/// one portrait, the same weights. A frame is fitted into whichever
/// matches its aspect, so a portrait set does not spend half the
/// detector's width on padding — which is what the 6D fixture did before
/// the second export existed (512 × 768 of a 1024 × 768 input). A
/// different size is a different file (`tools/export-xfeat.sh`).
pub const INPUT_LANDSCAPE: (usize, usize) = (1024, 768);
pub const INPUT_PORTRAIT: (usize, usize) = (768, 1024);

/// The long edge of the detector's input, for callers sizing a proxy.
pub const INPUT_LONG_EDGE: usize = 1024;

#[cfg(feature = "embedded-model")]
const EMBEDDED_LANDSCAPE: &[u8] = include_bytes!("../../../models/keypoints/xfeat-1024.onnx");
#[cfg(feature = "embedded-model")]
const EMBEDDED_PORTRAIT: &[u8] = include_bytes!("../../../models/keypoints/xfeat-768.onnx");

/// A loaded detector: the network at both shapes.
pub struct XFeat {
    landscape: ort::session::Session,
    portrait: ort::session::Session,
    pub options: DecodeOptions,
}

impl XFeat {
    /// The weights compiled into the binary.
    #[cfg(feature = "embedded-model")]
    pub fn embedded() -> Result<Self, PanoError> {
        Self::from_bytes(EMBEDDED_LANDSCAPE, EMBEDDED_PORTRAIT)
    }

    /// From the two exports on disk.
    pub fn from_paths(landscape: &std::path::Path, portrait: &std::path::Path) -> Result<Self, PanoError> {
        let l = std::fs::read(landscape).map_err(PanoError::ModelRead)?;
        let p = std::fs::read(portrait).map_err(PanoError::ModelRead)?;
        Self::from_bytes(&l, &p)
    }

    pub fn from_bytes(landscape: &[u8], portrait: &[u8]) -> Result<Self, PanoError> {
        install_backend();
        let session = |bytes: &[u8]| {
            ort::session::Session::builder()
                .map_err(PanoError::Inference)?
                .commit_from_memory(bytes)
                .map_err(PanoError::Inference)
        };
        Ok(XFeat {
            landscape: session(landscape)?,
            portrait: session(portrait)?,
            options: DecodeOptions::default(),
        })
    }

    /// Detect keypoints in an upright grayscale image.
    ///
    /// The image is fitted into the network's input of matching aspect —
    /// scaled down if larger, never up, and padded to the right and bottom
    /// — and the keypoints come back in the coordinates of `image` itself,
    /// so a caller that already scaled a frame to a proxy maps them on with
    /// the scale it used and nothing else.
    pub fn detect(&mut self, image: &Gray) -> Result<Features, PanoError> {
        let ((in_w, in_h), session) = if image.height > image.width {
            (INPUT_PORTRAIT, &mut self.portrait)
        } else {
            (INPUT_LANDSCAPE, &mut self.landscape)
        };
        let (fitted, scale) = image.fitted(in_w, in_h);
        let padded = fitted.padded(in_w, in_h);

        let input = ndarray::Array::from_shape_vec(ndarray::IxDyn(&[1, 1, in_h, in_w]), padded.data)
            .expect("shape matches the buffer by construction");
        let tensor = ort::value::Tensor::from_array(input).map_err(PanoError::Inference)?;
        let outputs = session
            .run(ort::inputs![tensor])
            .map_err(PanoError::Inference)?;

        let (w8, h8) = (in_w / 8, in_h / 8);
        let expect = |i: usize, channels: usize| -> Result<Vec<f32>, PanoError> {
            let (shape, data) = outputs[i]
                .try_extract_tensor::<f32>()
                .map_err(PanoError::Inference)?;
            let dims: Vec<i64> = shape.iter().copied().collect();
            if dims != [1, channels as i64, h8 as i64, w8 as i64] {
                return Err(PanoError::Model(format!(
                    "output {i} is {dims:?}, expected [1, {channels}, {h8}, {w8}] — \
                     not the export this decoder was written for"
                )));
            }
            Ok(data.to_vec())
        };
        let feats = expect(0, DESCRIPTOR_LEN)?;
        let keypoints = expect(1, 65)?;
        let heatmap = expect(2, 1)?;

        let mut features = decode_xfeat(
            &XFeatMaps {
                feats: &feats,
                keypoints: &keypoints,
                heatmap: &heatmap,
                width: w8,
                height: h8,
            },
            &self.options,
        );

        // Back to the caller's image: drop anything the padding produced,
        // undo the fit.
        let border = self.options.border as f32;
        let limit_x = fitted.width as f32 - border;
        let limit_y = fitted.height as f32 - border;
        let mut kept_kp = Vec::with_capacity(features.len());
        let mut kept_desc = Vec::with_capacity(features.descriptors.len());
        for (i, kp) in features.keypoints.iter().enumerate() {
            if kp.x >= limit_x || kp.y >= limit_y {
                continue;
            }
            kept_kp.push(crate::features::Keypoint {
                x: (kp.x / scale as f32),
                y: (kp.y / scale as f32),
                score: kp.score,
            });
            kept_desc.extend_from_slice(features.descriptor(i));
        }
        features.keypoints = kept_kp;
        features.descriptors = kept_desc;
        features.width = image.width;
        features.height = image.height;
        Ok(features)
    }
}

fn install_backend() {
    use std::sync::Once;
    static ONCE: Once = Once::new();
    ONCE.call_once(|| {
        // False if another crate installed it first, which is fine: there is
        // one backend compiled in for it to have chosen.
        let _ = ort::set_api(ort_tract::api());
    });
}