wifi-densepose-train 0.3.1

//! Cross-domain evaluation metrics (MERIDIAN Phase 6).
//!
//! MPJPE, domain gap ratio, and adaptation speedup for measuring how well a
//! WiFi-DensePose model generalizes across environments and hardware.

use std::collections::HashMap;

/// Aggregated cross-domain evaluation metrics.
#[derive(Debug, Clone)]
pub struct CrossDomainMetrics {
    /// In-domain (source) MPJPE (mm).
    pub in_domain_mpjpe: f32,
    /// Cross-domain (unseen environment) MPJPE (mm).
    pub cross_domain_mpjpe: f32,
    /// MPJPE after few-shot adaptation (mm).
    pub few_shot_mpjpe: f32,
    /// MPJPE across different WiFi hardware (mm).
    pub cross_hardware_mpjpe: f32,
    /// cross-domain / in-domain MPJPE. Target: < 1.5.
    pub domain_gap_ratio: f32,
    /// Labelled-sample savings vs training from scratch.
    pub adaptation_speedup: f32,
}

/// Evaluates pose estimation across multiple domains.
///
/// Domain 0 = in-domain (source); other IDs = cross-domain.
///
/// ```rust
/// use wifi_densepose_train::eval::{CrossDomainEvaluator, mpjpe};
/// let ev = CrossDomainEvaluator::new(17);
/// let preds = vec![(vec![0.0_f32; 51], vec![0.0_f32; 51])];
/// let m = ev.evaluate(&preds, &[0]);
/// assert!(m.in_domain_mpjpe >= 0.0);
/// ```
pub struct CrossDomainEvaluator {
    n_joints: usize,
}

impl CrossDomainEvaluator {
    /// Create evaluator for `n_joints` body joints (e.g. 17 for COCO).
    pub fn new(n_joints: usize) -> Self { Self { n_joints } }

    /// Evaluate predictions grouped by domain. Each pair is (predicted, gt)
    /// with `n_joints * 3` floats. `domain_labels` must match length.
    pub fn evaluate(&self, predictions: &[(Vec<f32>, Vec<f32>)], domain_labels: &[u32]) -> CrossDomainMetrics {
        assert_eq!(predictions.len(), domain_labels.len(), "length mismatch");
        let mut by_dom: HashMap<u32, Vec<f32>> = HashMap::new();
        for (i, (p, g)) in predictions.iter().enumerate() {
            by_dom.entry(domain_labels[i]).or_default().push(mpjpe(p, g, self.n_joints));
        }
        let in_dom = mean_of(by_dom.get(&0));
        let cross_errs: Vec<f32> = by_dom.iter().filter(|(&d, _)| d != 0).flat_map(|(_, e)| e.iter().copied()).collect();
        let cross_dom = if cross_errs.is_empty() { 0.0 } else { cross_errs.iter().sum::<f32>() / cross_errs.len() as f32 };
        let few_shot = if by_dom.contains_key(&2) { mean_of(by_dom.get(&2)) } else { (in_dom + cross_dom) / 2.0 };
        let cross_hw = if by_dom.contains_key(&3) { mean_of(by_dom.get(&3)) } else { cross_dom };
        let gap = if in_dom > 1e-10 { cross_dom / in_dom } else if cross_dom > 1e-10 { f32::INFINITY } else { 1.0 };
        let speedup = if few_shot > 1e-10 { cross_dom / few_shot } else { 1.0 };
        CrossDomainMetrics { in_domain_mpjpe: in_dom, cross_domain_mpjpe: cross_dom, few_shot_mpjpe: few_shot,
            cross_hardware_mpjpe: cross_hw, domain_gap_ratio: gap, adaptation_speedup: speedup }
    }
}

/// Mean Per Joint Position Error: average Euclidean distance across `n_joints`.
///
/// `pred` and `gt` are flat `[n_joints * 3]` (x, y, z per joint).
pub fn mpjpe(pred: &[f32], gt: &[f32], n_joints: usize) -> f32 {
    if n_joints == 0 { return 0.0; }
    let total: f32 = (0..n_joints).map(|j| {
        let b = j * 3;
        let d = |off| pred.get(b + off).copied().unwrap_or(0.0) - gt.get(b + off).copied().unwrap_or(0.0);
        (d(0).powi(2) + d(1).powi(2) + d(2).powi(2)).sqrt()
    }).sum();
    total / n_joints as f32
}

fn mean_of(v: Option<&Vec<f32>>) -> f32 {
    match v { Some(e) if !e.is_empty() => e.iter().sum::<f32>() / e.len() as f32, _ => 0.0 }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn mpjpe_known_value() {
        assert!((mpjpe(&[0.0, 0.0, 0.0], &[3.0, 4.0, 0.0], 1) - 5.0).abs() < 1e-6);
    }

    #[test]
    fn mpjpe_two_joints() {
        // Joint 0: dist=5, Joint 1: dist=0 -> mean=2.5
        assert!((mpjpe(&[0.0,0.0,0.0, 1.0,1.0,1.0], &[3.0,4.0,0.0, 1.0,1.0,1.0], 2) - 2.5).abs() < 1e-6);
    }

    #[test]
    fn mpjpe_zero_when_identical() {
        let c = vec![1.5, 2.3, 0.7, 4.1, 5.9, 3.2];
        assert!(mpjpe(&c, &c, 2).abs() < 1e-10);
    }

    #[test]
    fn mpjpe_zero_joints() { assert_eq!(mpjpe(&[], &[], 0), 0.0); }

    #[test]
    fn domain_gap_ratio_computed() {
        let ev = CrossDomainEvaluator::new(1);
        let preds = vec![
            (vec![0.0,0.0,0.0], vec![1.0,0.0,0.0]), // dom 0, err=1
            (vec![0.0,0.0,0.0], vec![2.0,0.0,0.0]), // dom 1, err=2
        ];
        let m = ev.evaluate(&preds, &[0, 1]);
        assert!((m.in_domain_mpjpe - 1.0).abs() < 1e-6);
        assert!((m.cross_domain_mpjpe - 2.0).abs() < 1e-6);
        assert!((m.domain_gap_ratio - 2.0).abs() < 1e-6);
    }

    #[test]
    fn evaluate_groups_by_domain() {
        let ev = CrossDomainEvaluator::new(1);
        let preds = vec![
            (vec![0.0,0.0,0.0], vec![1.0,0.0,0.0]),
            (vec![0.0,0.0,0.0], vec![3.0,0.0,0.0]),
            (vec![0.0,0.0,0.0], vec![5.0,0.0,0.0]),
        ];
        let m = ev.evaluate(&preds, &[0, 0, 1]);
        assert!((m.in_domain_mpjpe - 2.0).abs() < 1e-6);
        assert!((m.cross_domain_mpjpe - 5.0).abs() < 1e-6);
    }

    #[test]
    fn domain_gap_perfect() {
        let ev = CrossDomainEvaluator::new(1);
        let preds = vec![(vec![1.0,2.0,3.0], vec![1.0,2.0,3.0]), (vec![4.0,5.0,6.0], vec![4.0,5.0,6.0])];
        assert!((ev.evaluate(&preds, &[0, 1]).domain_gap_ratio - 1.0).abs() < 1e-6);
    }

    #[test]
    fn evaluate_multiple_cross_domains() {
        let ev = CrossDomainEvaluator::new(1);
        let preds = vec![
            (vec![0.0,0.0,0.0], vec![1.0,0.0,0.0]),
            (vec![0.0,0.0,0.0], vec![4.0,0.0,0.0]),
            (vec![0.0,0.0,0.0], vec![6.0,0.0,0.0]),
        ];
        let m = ev.evaluate(&preds, &[0, 1, 3]);
        assert!((m.in_domain_mpjpe - 1.0).abs() < 1e-6);
        assert!((m.cross_domain_mpjpe - 5.0).abs() < 1e-6);
        assert!((m.cross_hardware_mpjpe - 6.0).abs() < 1e-6);
    }
}