use crate::image::Image;
use burn::tensor::backend::Backend;
use std::time::Duration;
#[derive(Clone, Debug)]
pub struct Frame<B: Backend> {
pub image: Image<B>,
pub pts: Duration,
pub duration: Duration,
pub index: usize,
pub is_keyframe: bool,
}
impl<B: Backend> Frame<B> {
#[must_use]
pub fn new(image: Image<B>, pts: Duration, index: usize) -> Self {
Self {
image,
pts,
duration: Duration::ZERO,
index,
is_keyframe: false,
}
}
#[must_use]
pub fn keyframe(image: Image<B>, pts: Duration, index: usize) -> Self {
Self {
image,
pts,
duration: Duration::ZERO,
index,
is_keyframe: true,
}
}
#[must_use]
pub fn with_duration(mut self, duration: Duration) -> Self {
self.duration = duration;
self
}
#[must_use]
pub fn width(&self) -> usize {
self.image.width()
}
#[must_use]
pub fn height(&self) -> usize {
self.image.height()
}
#[must_use]
pub fn channels(&self) -> usize {
self.image.channels()
}
#[must_use]
pub fn shape(&self) -> [usize; 3] {
self.image.shape()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::{TestBackend, test_device};
use burn::tensor::{Tensor, TensorData};
#[test]
fn test_frame_creation() {
let device = test_device();
let data = TensorData::new(vec![0.5f32; 3 * 64 * 64], [3, 64, 64]);
let tensor = Tensor::<TestBackend, 3>::from_data(data, &device);
let img = Image::new(tensor);
let frame = Frame::new(img, Duration::from_millis(33), 0);
assert_eq!(frame.width(), 64);
assert_eq!(frame.height(), 64);
assert_eq!(frame.channels(), 3);
assert_eq!(frame.index, 0);
assert!(!frame.is_keyframe);
let kf = Frame::keyframe(frame.image.clone(), Duration::from_millis(0), 0);
assert!(kf.is_keyframe);
}
#[test]
fn test_frame_with_duration() {
let device = test_device();
let data = TensorData::new(vec![0.5f32; 3 * 32 * 32], [3, 32, 32]);
let tensor = Tensor::<TestBackend, 3>::from_data(data, &device);
let img = Image::new(tensor);
let frame =
Frame::new(img, Duration::from_millis(0), 0).with_duration(Duration::from_millis(33));
assert_eq!(frame.duration, Duration::from_millis(33));
}
}