eventcv_core/
interp.rs

1//! Learned frame interpolation, run *before* simulation.
2//!
3//! [`crate::simulate::Upsample`] already subdivides the interval between two frames, but the levels
4//! it subdivides are a **linear** blend of them. That is exactly right when the intensity at a pixel
5//! moves linearly over a frame gap, and wrong when it does not — an edge crossing a pixel makes it
6//! step, not ramp, and a linear blend puts every event it should have produced at the wrong moment.
7//! v2e reaches for Super-SloMo here; this reaches for whatever ONNX graph the caller has.
8//!
9//! It sits *before* the simulator rather than inside it, which is also where v2e puts it: the
10//! interpolated frames are simply more source frames, with timestamps to match, and the simulator's
11//! own [`Upsample`](crate::simulate::Upsample) then refines what is left. That means the pixel model
12//! and its hot loop are untouched, and a run without an interpolator is byte-for-byte the run that
13//! came before this module existed.
14//!
15//! No weights are bundled. eventcv runs graphs; it does not ship a zoo. Export RIFE (or anything
16//! with the same shape) yourself and hand over the path.
17
18use crate::representation::RepresentationError;
19
20/// Produces frames between two others.
21///
22/// Frames are single-plane luma in `[0, 1]`, row-major — what [`crate::simulate::Simulator`]
23/// consumes and what the video decoder already reduces its RGB to. A model trained on colour still
24/// works: a grey image is a valid one, and the simulator would have thrown the colour away anyway.
25pub trait FrameInterpolator: Send {
26    /// The frames at each of `fractions` (each strictly between 0 and 1) along the path from `a`
27    /// to `b`, in the order given.
28    fn between(
29        &mut self,
30        a: &[f32],
31        b: &[f32],
32        width: usize,
33        height: usize,
34        fractions: &[f32],
35    ) -> Result<Vec<Vec<f32>>, InterpError>;
36}
37
38#[derive(Debug)]
39pub enum InterpError {
40    /// The graph's inputs are not a shape this can drive. Carries what it saw and what it wanted.
41    Unsupported(String),
42    /// The model itself failed.
43    Model(String),
44}
45
46impl std::fmt::Display for InterpError {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        match self {
49            Self::Unsupported(message) | Self::Model(message) => formatter.write_str(message),
50        }
51    }
52}
53
54impl std::error::Error for InterpError {}
55
56impl From<InterpError> for RepresentationError {
57    fn from(error: InterpError) -> Self {
58        Self::Device(error.to_string())
59    }
60}
61
62/// The linear blend the simulator already does, as a [`FrameInterpolator`].
63///
64/// Not useful in production — the simulator interpolates linearly on its own, for free — but it is
65/// what makes the *plumbing* testable without an ONNX export to hand, and it is the baseline any
66/// learned interpolator has to beat.
67pub struct LinearInterpolator;
68
69impl FrameInterpolator for LinearInterpolator {
70    fn between(
71        &mut self,
72        a: &[f32],
73        b: &[f32],
74        _width: usize,
75        _height: usize,
76        fractions: &[f32],
77    ) -> Result<Vec<Vec<f32>>, InterpError> {
78        Ok(fractions
79            .iter()
80            .map(|fraction| {
81                a.iter()
82                    .zip(b)
83                    .map(|(a, b)| a + (b - a) * fraction)
84                    .collect()
85            })
86            .collect())
87    }
88}
89
90/// How an ONNX interpolator wants its frame pair.
91///
92/// Exports disagree, and the disagreement is entirely about packaging rather than about what the
93/// network does, so it is detected from the declared inputs rather than configured.
94#[cfg(feature = "onnx")]
95#[derive(Clone, Debug, PartialEq, Eq)]
96enum Layout {
97    /// One input of six channels: the two frames stacked.
98    Stacked { image: String, timestep: Option<String> },
99    /// Two image inputs.
100    Separate {
101        first: String,
102        second: String,
103        timestep: Option<String>,
104    },
105}
106
107/// Drives an ONNX frame interpolator — RIFE and anything shaped like it.
108///
109/// # Which exports work
110///
111/// The graph must take a frame pair, either stacked into one six-channel input or as two
112/// three-channel ones, and return one image. A scalar `timestep` input (RIFE v4 and later) is used
113/// when present, which is what allows an arbitrary fraction in one call. Without it the network can
114/// only produce the midpoint, so a fraction of `k/2^n` is reached by bisecting `n` times and
115/// anything else is refused rather than approximated.
116#[cfg(feature = "onnx")]
117pub struct OnnxInterpolator {
118    model: crate::model::Model,
119    layout: Layout,
120}
121
122#[cfg(feature = "onnx")]
123impl OnnxInterpolator {
124    /// Loads `path` and works out how to feed it.
125    pub fn load(path: &str) -> Result<Self, InterpError> {
126        let model = crate::model::Model::load(path).map_err(|error| InterpError::Model(error.to_string()))?;
127        let layout = Self::detect(&model)?;
128        Ok(Self { model, layout })
129    }
130
131    /// Reads the graph's declared inputs and decides how the pair is passed.
132    fn detect(model: &crate::model::Model) -> Result<Layout, InterpError> {
133        // A timestep is the input that is not an image: rank below three, or a single element.
134        let images: Vec<&crate::model::Port> = model
135            .inputs()
136            .iter()
137            .filter(|port| port.shape.len() >= 3)
138            .collect();
139        let timestep = model
140            .inputs()
141            .iter()
142            .find(|port| port.shape.len() < 3 || port.shape.iter().all(|dim| *dim == 1))
143            .map(|port| port.name.clone());
144        let channels = |port: &crate::model::Port| port.shape.get(port.shape.len() - 3).copied();
145
146        match images.as_slice() {
147            [image] if matches!(channels(image), Some(6) | Some(-1)) => Ok(Layout::Stacked {
148                image: image.name.clone(),
149                timestep,
150            }),
151            [first, second] => Ok(Layout::Separate {
152                first: first.name.clone(),
153                second: second.name.clone(),
154                timestep,
155            }),
156            other => Err(InterpError::Unsupported(format!(
157                "a frame interpolator should take a frame pair — one six-channel input or two \
158                 three-channel ones, optionally with a timestep — but this graph declares {} \
159                 image-shaped inputs: {:?}",
160                other.len(),
161                model
162                    .inputs()
163                    .iter()
164                    .map(|port| (port.name.as_str(), &port.shape))
165                    .collect::<Vec<_>>()
166            ))),
167        }
168    }
169
170    fn timestep(&self) -> Option<&str> {
171        match &self.layout {
172            Layout::Stacked { timestep, .. } | Layout::Separate { timestep, .. } => {
173                timestep.as_deref()
174            }
175        }
176    }
177
178    /// One forward pass: the frame `fraction` of the way from `a` to `b`.
179    fn once(
180        &mut self,
181        a: &[f32],
182        b: &[f32],
183        width: usize,
184        height: usize,
185        fraction: f32,
186    ) -> Result<Vec<f32>, InterpError> {
187        use ndarray::{Array, IxDyn};
188
189        // Luma replicated across RGB: these networks are trained on colour, and a grey image is
190        // simply one where the channels agree.
191        let rgb = |plane: &[f32]| -> Vec<f32> {
192            let mut out = Vec::with_capacity(plane.len() * 3);
193            for _ in 0..3 {
194                out.extend_from_slice(plane);
195            }
196            out
197        };
198        let shaped = |data: Vec<f32>, channels: usize| {
199            Array::from_shape_vec(IxDyn(&[1, channels, height, width]), data)
200                .map_err(|error| InterpError::Model(error.to_string()))
201        };
202
203        let mut inputs = match self.layout.clone() {
204            Layout::Stacked { image, .. } => {
205                let mut both = rgb(a);
206                both.extend(rgb(b));
207                vec![(image, shaped(both, 6)?)]
208            }
209            Layout::Separate { first, second, .. } => vec![
210                (first, shaped(rgb(a), 3)?),
211                (second, shaped(rgb(b), 3)?),
212            ],
213        };
214        if let Some(name) = self.timestep().map(str::to_owned) {
215            inputs.push((
216                name,
217                Array::from_shape_vec(IxDyn(&[1]), vec![fraction])
218                    .map_err(|error| InterpError::Model(error.to_string()))?,
219            ));
220        }
221
222        let outputs = self
223            .model
224            .run_named(inputs)
225            .map_err(|error| InterpError::Model(error.to_string()))?;
226        let (_, image) = outputs
227            .into_iter()
228            .next()
229            .ok_or_else(|| InterpError::Model("the graph returned nothing".into()))?;
230        Ok(to_luma(image.as_slice().unwrap_or(&[]), width * height))
231    }
232
233    /// The frame at `fraction` when the graph has no timestep input: bisect towards it, which only
234    /// reaches fractions of the form `k / 2^n`.
235    fn bisect(
236        &mut self,
237        a: &[f32],
238        b: &[f32],
239        width: usize,
240        height: usize,
241        fraction: f32,
242        depth: usize,
243    ) -> Result<Vec<f32>, InterpError> {
244        const MAX_DEPTH: usize = 4;
245        if (fraction - 0.5).abs() < 1e-6 {
246            return self.once(a, b, width, height, 0.5);
247        }
248        if depth >= MAX_DEPTH {
249            return Err(InterpError::Unsupported(format!(
250                "this export has no timestep input, so it can only produce midpoints; a fraction \
251                 of {fraction} would need more than {MAX_DEPTH} bisections. Use an interpolation \
252                 factor that is a power of two, or export a model that takes a timestep."
253            )));
254        }
255        let middle = self.once(a, b, width, height, 0.5)?;
256        if fraction < 0.5 {
257            self.bisect(a, &middle, width, height, fraction * 2.0, depth + 1)
258        } else {
259            self.bisect(&middle, b, width, height, (fraction - 0.5) * 2.0, depth + 1)
260        }
261    }
262}
263
264/// Collapses a model's `[1, C, H, W]` output back to one luma plane.
265#[cfg(feature = "onnx")]
266fn to_luma(data: &[f32], plane: usize) -> Vec<f32> {
267    if plane == 0 {
268        return Vec::new();
269    }
270    let channels = (data.len() / plane).max(1);
271    (0..plane)
272        .map(|index| {
273            let sum: f32 = (0..channels.min(3))
274                .map(|channel| data.get(channel * plane + index).copied().unwrap_or(0.0))
275                .sum();
276            (sum / channels.min(3) as f32).clamp(0.0, 1.0)
277        })
278        .collect()
279}
280
281#[cfg(feature = "onnx")]
282impl FrameInterpolator for OnnxInterpolator {
283    fn between(
284        &mut self,
285        a: &[f32],
286        b: &[f32],
287        width: usize,
288        height: usize,
289        fractions: &[f32],
290    ) -> Result<Vec<Vec<f32>>, InterpError> {
291        let timed = self.timestep().is_some();
292        fractions
293            .iter()
294            .map(|fraction| {
295                if timed {
296                    self.once(a, b, width, height, *fraction)
297                } else {
298                    self.bisect(a, b, width, height, *fraction, 0)
299                }
300            })
301            .collect()
302    }
303}
304
305/// An interpolator and how many sub-frames it should produce.
306///
307/// `factor` is the number of intervals each source pair becomes, so `4` inserts three frames.
308pub struct Interpolation<'a> {
309    pub interpolator: &'a mut dyn FrameInterpolator,
310    pub factor: usize,
311}
312
313impl Interpolation<'_> {
314    /// The fractions between two source frames, `factor - 1` of them.
315    pub fn fractions(&self) -> Vec<f32> {
316        (1..self.factor.max(1))
317            .map(|step| step as f32 / self.factor as f32)
318            .collect()
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::{FrameInterpolator, Interpolation, LinearInterpolator};
325
326    #[test]
327    fn fractions_split_the_interval_evenly_and_exclude_the_endpoints() {
328        let mut linear = LinearInterpolator;
329        let plan = Interpolation {
330            interpolator: &mut linear,
331            factor: 4,
332        };
333        assert_eq!(plan.fractions(), vec![0.25, 0.5, 0.75]);
334    }
335
336    #[test]
337    fn a_factor_of_one_inserts_nothing() {
338        let mut linear = LinearInterpolator;
339        let plan = Interpolation {
340            interpolator: &mut linear,
341            factor: 1,
342        };
343        assert!(plan.fractions().is_empty());
344    }
345
346    #[test]
347    fn the_linear_baseline_blends_the_way_the_simulator_would() {
348        let mut linear = LinearInterpolator;
349        let frames = linear
350            .between(&[0.0, 1.0], &[1.0, 0.0], 2, 1, &[0.25, 0.75])
351            .unwrap();
352        assert_eq!(frames, vec![vec![0.25, 0.75], vec![0.75, 0.25]]);
353    }
354}
eventcv_core/interp.rs

eventcv_core/
interp.rs
