Refactor video rendering code

Author: moonfloof

src/data/mod.rs

@@ -1,3 +1,4 @@
 pub mod cli;
 pub mod song;
 pub mod track;
+pub mod video;

src/data/song.rs

@@ -1,21 +1,18 @@
-use super::track::load_track_into_memory;
+use super::{track::load_track_into_memory, video::Encoding};
 use crate::{
 	display::{draw, RGB},
-	SCREEN_FRAME_RATE, SCREEN_HEIGHT, SCREEN_SCALE, SCREEN_WIDTH,
+	SCREEN_FRAME_RATE, SCREEN_HEIGHT, SCREEN_WIDTH,
 };
-use fast_image_resize as fr;
 use image::RgbImage;
-use ndarray::Array3;
 use rayon::prelude::*;
 use serde::Deserialize;
-use std::{collections::VecDeque, fs::File, io::BufReader, num::NonZeroU32};
+use std::{collections::VecDeque, fs::File, io::BufReader};
 use symphonia::core::{
 	audio::{AudioBufferRef, Signal},
 	codecs::Decoder,
 	errors::Error,
 	formats::{FormatReader, Track},
 };
-use video_rs::{Encoder, Time};
 
 #[derive(Deserialize)]
 pub struct Channel {
@@ -162,14 +159,6 @@ impl Window {
 	}
 }
 
-pub struct Encoding {
-	pub encoder: Encoder,
-	pub position: Time,
-	pub resizer: fr::Resizer,
-	pub size_src: (NonZeroU32, NonZeroU32),
-	pub size_dst: (NonZeroU32, NonZeroU32),
-}
-
 #[derive(Deserialize)]
 pub struct Song {
 	pub channels: Vec<Channel>,
@@ -320,41 +309,7 @@ impl Song {
 		}
 
 		// Render frame to video
-
-		let width = *SCREEN_WIDTH * *SCREEN_SCALE;
-		let height = *SCREEN_HEIGHT * *SCREEN_SCALE;
-
-		let rs: Vec<u8> = if *SCREEN_SCALE == 1 {
-			frame.as_raw().to_vec()
-		} else {
-			let src_image = fr::Image::from_slice_u8(
-				encoding.size_src.0,
-				encoding.size_src.1,
-				frame,
-				fr::PixelType::U8x3,
-			)
-			.unwrap();
-
-			let mut dst_image = fr::Image::new(
-				encoding.size_dst.0,
-				encoding.size_dst.1,
-				fr::PixelType::U8x3,
-			);
-
-			// Get mutable view of destination image data
-			let mut dst_view = dst_image.view_mut();
-			encoding
-				.resizer
-				.resize(&src_image.view(), &mut dst_view)
-				.unwrap();
-
-			dst_image.buffer().to_vec()
-		};
-
-		let ef: Array3<u8> =
-			ndarray::Array3::from_shape_vec((height as usize, width as usize, 3), rs).unwrap();
-
-		encoding.encoder.encode(&ef, &encoding.position).unwrap();
+		encoding.render_frame(frame);
 
 		self.frame += 1;
 		if self.frame % 100 == 0 {

src/data/video.rs

@@ -0,0 +1,78 @@
+use crate::{SCREEN_FRAME_RATE, SCREEN_HEIGHT, SCREEN_SCALE, SCREEN_WIDTH};
+use fast_image_resize as fr;
+use image::RgbImage;
+use ndarray::Array3;
+use std::num::NonZeroU32;
+use video_rs::{Encoder, Time};
+
+pub struct Encoding {
+	pub encoder: Encoder,
+	pub position: Time,
+	pub frame_duration: Time,
+	pub resizer: fr::Resizer,
+	pub size_src: (NonZeroU32, NonZeroU32),
+	pub size_dst: (NonZeroU32, NonZeroU32),
+}
+
+impl Encoding {
+	pub fn new(encoder: Encoder) -> Self {
+		Encoding {
+			encoder,
+			position: Time::zero(),
+			frame_duration: Time::from_nth_of_a_second(*SCREEN_FRAME_RATE),
+			size_src: (
+				NonZeroU32::new(*SCREEN_WIDTH).unwrap(),
+				NonZeroU32::new(*SCREEN_HEIGHT).unwrap(),
+			),
+			size_dst: (
+				NonZeroU32::new(*SCREEN_WIDTH * *SCREEN_SCALE).unwrap(),
+				NonZeroU32::new(*SCREEN_HEIGHT * *SCREEN_SCALE).unwrap(),
+			),
+			resizer: fr::Resizer::new(fr::ResizeAlg::Nearest),
+		}
+	}
+
+	pub fn resize_frame(&mut self, frame: &mut RgbImage) -> Vec<u8> {
+		if *SCREEN_SCALE == 1 {
+			return frame.as_raw().to_vec();
+		}
+
+		let src_image =
+			fr::Image::from_slice_u8(self.size_src.0, self.size_src.1, frame, fr::PixelType::U8x3)
+				.unwrap();
+
+		let mut dst_image = fr::Image::new(self.size_dst.0, self.size_dst.1, fr::PixelType::U8x3);
+
+		// Get mutable view of destination image data
+		let mut dst_view = dst_image.view_mut();
+		self.resizer
+			.resize(&src_image.view(), &mut dst_view)
+			.unwrap();
+
+		dst_image.buffer().to_vec()
+	}
+
+	pub fn render_frame(&mut self, frame: &mut RgbImage) {
+		let pixels = self.resize_frame(frame);
+		let ef: Array3<u8> = ndarray::Array3::from_shape_vec(
+			(
+				self.size_dst.1.get() as usize,
+				self.size_dst.0.get() as usize,
+				3,
+			),
+			pixels,
+		)
+		.unwrap();
+
+		self.encoder.encode(&ef, &self.position).unwrap();
+		self.update_position();
+	}
+
+	pub fn update_position(&mut self) {
+		self.position = self.position.aligned_with(&self.frame_duration).add();
+	}
+
+	pub fn flush(&mut self) {
+		self.encoder.finish().unwrap();
+	}
+}

src/main.rs

@@ -1,11 +1,10 @@
-use crate::data::{cli::Args, song::Encoding};
+use crate::data::{cli::Args, video::Encoding};
 use clap::Parser;
 use data::song::{Song, Window};
-use fast_image_resize as fr;
 use image::RgbImage;
 use lazy_static::lazy_static;
-use std::{num::NonZeroU32, path::PathBuf};
-use video_rs::{Encoder, EncoderSettings, Locator, Time};
+use std::path::PathBuf;
+use video_rs::{Encoder, EncoderSettings, Locator};
 
 mod data;
 mod display;
@@ -32,30 +31,16 @@ fn main() {
 
 	let destination: Locator = PathBuf::from(&song.video_file_out).into();
 	let settings = EncoderSettings::for_h264_yuv420p(width, height, false);
-	let duration: Time = Time::from_nth_of_a_second(*SCREEN_FRAME_RATE);
-
 	let mut frame = RgbImage::new(*SCREEN_WIDTH, *SCREEN_HEIGHT);
+	let mut encoding =
+		Encoding::new(Encoder::new(&destination, settings).expect("failed to create encoder"));
 
 	println!(
 		"\nGenerated frame of size {}x{}",
 		frame.width(),
 		frame.height()
 	);
 
-	let mut encoding = Encoding {
-		encoder: Encoder::new(&destination, settings).expect("failed to create encoder"),
-		position: Time::zero(),
-		size_src: (
-			NonZeroU32::new(*SCREEN_WIDTH).unwrap(),
-			NonZeroU32::new(*SCREEN_HEIGHT).unwrap(),
-		),
-		size_dst: (
-			NonZeroU32::new(width as u32).unwrap(),
-			NonZeroU32::new(height as u32).unwrap(),
-		),
-		resizer: fr::Resizer::new(fr::ResizeAlg::Nearest),
-	};
-
 	// Step 2: Render waveforms
 	loop {
 		let result = song.draw(&mut frame, &mut encoding);
@@ -64,10 +49,8 @@ fn main() {
 			println!("{:?}", result.err().unwrap());
 			break;
 		}
-
-		encoding.position = encoding.position.aligned_with(&duration).add();
 	}
 
 	// Step 3: Flush MP4 to file
-	encoding.encoder.finish().unwrap();
+	encoding.flush();
 }