Skip to main content

oximedia_optimize/reference/
selection.rs

1//! Reference frame selection optimization.
2
3/// Reference frame score.
4#[derive(Debug, Clone, Copy)]
5pub struct RefFrameScore {
6    /// Frame index.
7    pub frame_idx: usize,
8    /// Prediction quality score.
9    pub quality_score: f64,
10    /// Temporal distance.
11    pub temporal_distance: i32,
12    /// Total score (higher is better).
13    pub total_score: f64,
14}
15
16impl RefFrameScore {
17    /// Creates a new reference frame score.
18    #[must_use]
19    pub fn new(frame_idx: usize, quality_score: f64, temporal_distance: i32) -> Self {
20        // Prefer closer frames and higher quality
21        let distance_penalty = f64::from(temporal_distance.abs()) * 0.1;
22        let total_score = quality_score - distance_penalty;
23
24        Self {
25            frame_idx,
26            quality_score,
27            temporal_distance,
28            total_score,
29        }
30    }
31}
32
33/// Reference selection optimizer.
34pub struct ReferenceSelection {
35    max_references: usize,
36    #[allow(dead_code)]
37    temporal_bias: f64,
38}
39
40impl Default for ReferenceSelection {
41    fn default() -> Self {
42        Self::new(3, 0.1)
43    }
44}
45
46impl ReferenceSelection {
47    /// Creates a new reference selector.
48    #[must_use]
49    pub fn new(max_references: usize, temporal_bias: f64) -> Self {
50        Self {
51            max_references,
52            temporal_bias,
53        }
54    }
55
56    /// Selects best reference frames.
57    #[allow(dead_code)]
58    #[must_use]
59    pub fn select_references(
60        &self,
61        current_poc: i32,
62        available_frames: &[(usize, i32)], // (index, POC)
63        src: &[u8],
64    ) -> Vec<usize> {
65        let mut scores: Vec<RefFrameScore> = available_frames
66            .iter()
67            .map(|&(idx, poc)| {
68                let temporal_distance = current_poc - poc;
69                let quality = self.estimate_quality(src, idx);
70                RefFrameScore::new(idx, quality, temporal_distance)
71            })
72            .collect();
73
74        // Sort by total score (descending)
75        scores.sort_by(|a, b| {
76            b.total_score
77                .partial_cmp(&a.total_score)
78                .unwrap_or(std::cmp::Ordering::Equal)
79        });
80
81        // Return top N references
82        scores
83            .iter()
84            .take(self.max_references)
85            .map(|s| s.frame_idx)
86            .collect()
87    }
88
89    fn estimate_quality(&self, _src: &[u8], _frame_idx: usize) -> f64 {
90        // Simplified quality estimation
91        // In production, would compare actual pixels
92        100.0
93    }
94
95    /// Determines if a frame should be kept as reference.
96    #[must_use]
97    pub fn should_keep_as_reference(&self, frame_type: FrameType, layer: u8) -> bool {
98        match frame_type {
99            FrameType::Key => true,         // Always keep key frames
100            FrameType::Inter => layer == 0, // Keep base layer inter frames
101        }
102    }
103}
104
105/// Frame types for reference management.
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
107pub enum FrameType {
108    /// Key frame (I-frame).
109    Key,
110    /// Inter frame (P/B-frame).
111    Inter,
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_ref_frame_score() {
120        let score = RefFrameScore::new(0, 100.0, 2);
121        assert_eq!(score.frame_idx, 0);
122        assert_eq!(score.quality_score, 100.0);
123        assert_eq!(score.temporal_distance, 2);
124        assert!(score.total_score < 100.0); // Has distance penalty
125    }
126
127    #[test]
128    fn test_reference_selection_creation() {
129        let selector = ReferenceSelection::default();
130        assert_eq!(selector.max_references, 3);
131    }
132
133    #[test]
134    fn test_select_references() {
135        let selector = ReferenceSelection::default();
136        let available = vec![(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)];
137        let src = vec![128u8; 64];
138        let selected = selector.select_references(5, &available, &src);
139        assert!(selected.len() <= 3);
140    }
141
142    #[test]
143    fn test_should_keep_reference() {
144        let selector = ReferenceSelection::default();
145        assert!(selector.should_keep_as_reference(FrameType::Key, 0));
146        assert!(selector.should_keep_as_reference(FrameType::Inter, 0));
147        assert!(!selector.should_keep_as_reference(FrameType::Inter, 1));
148    }
149}