warden_vision/core/
classification.rs

1//! Whole-image classification results, as opposed to the per-object
2//! detections in [`crate::core::Detections`].
3
4/// A single class prediction: class id plus confidence.
5#[derive(Debug, Clone, Copy, PartialEq)]
6pub struct ClassPrediction {
7    /// Index of the predicted class.
8    pub class_id: usize,
9    /// Confidence score, typically in `[0, 1]`.
10    pub confidence: f32,
11}
12
13impl ClassPrediction {
14    /// Creates a new class prediction.
15    pub fn new(class_id: usize, confidence: f32) -> Self {
16        Self {
17            class_id,
18            confidence,
19        }
20    }
21}
22
23/// The classification results for a single image (or crop): one
24/// [`ClassPrediction`] per class the classifier scored.
25#[derive(Debug, Clone, Default, PartialEq)]
26pub struct Classifications {
27    /// One entry per class, in whatever order the classifier produced
28    /// them (typically class id order, not sorted by confidence). Use
29    /// [`Classifications::top_k`] / [`Classifications::top1`] for a
30    /// confidence-sorted view.
31    pub predictions: Vec<ClassPrediction>,
32}
33
34impl Classifications {
35    /// Creates a new `Classifications` from an existing prediction list.
36    pub fn new(predictions: Vec<ClassPrediction>) -> Self {
37        Self { predictions }
38    }
39
40    /// Builds a `Classifications` from a dense per-class score vector
41    /// (`scores[class_id]`), as produced directly by a softmax/sigmoid
42    /// output layer.
43    pub fn from_scores(scores: &[f32]) -> Self {
44        Self::new(
45            scores
46                .iter()
47                .enumerate()
48                .map(|(class_id, &confidence)| ClassPrediction::new(class_id, confidence))
49                .collect(),
50        )
51    }
52
53    /// The `k` highest-confidence predictions, sorted descending. Returns
54    /// fewer than `k` if there aren't that many predictions.
55    pub fn top_k(&self, k: usize) -> Vec<ClassPrediction> {
56        let mut sorted = self.predictions.clone();
57        sorted.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
58        sorted.truncate(k);
59        sorted
60    }
61
62    /// The single highest-confidence prediction, if any.
63    pub fn top1(&self) -> Option<ClassPrediction> {
64        self.predictions
65            .iter()
66            .copied()
67            .max_by(|a, b| a.confidence.total_cmp(&b.confidence))
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn from_scores_indexes_predictions_by_position() {
77        let classifications = Classifications::from_scores(&[0.1, 0.7, 0.2]);
78        assert_eq!(classifications.predictions.len(), 3);
79        assert_eq!(classifications.predictions[1], ClassPrediction::new(1, 0.7));
80    }
81
82    #[test]
83    fn top1_returns_the_highest_confidence_prediction() {
84        let classifications = Classifications::from_scores(&[0.1, 0.7, 0.2]);
85        assert_eq!(classifications.top1(), Some(ClassPrediction::new(1, 0.7)));
86    }
87
88    #[test]
89    fn top1_of_empty_is_none() {
90        assert_eq!(Classifications::default().top1(), None);
91    }
92
93    #[test]
94    fn top_k_sorts_descending_and_truncates() {
95        let classifications = Classifications::from_scores(&[0.1, 0.7, 0.5, 0.2]);
96        let top2 = classifications.top_k(2);
97        assert_eq!(
98            top2,
99            vec![ClassPrediction::new(1, 0.7), ClassPrediction::new(2, 0.5)]
100        );
101    }
102
103    #[test]
104    fn top_k_larger_than_predictions_returns_all() {
105        let classifications = Classifications::from_scores(&[0.1, 0.7]);
106        assert_eq!(classifications.top_k(10).len(), 2);
107    }
108}
warden_vision/core/classification.rs

warden_vision/core/
classification.rs
