From e6a0429f39a86f895a829f9f4bd9aece8a06ec41 Mon Sep 17 00:00:00 2001
From: Janne Mareike Koschinski <janne@kuschku.de>
Date: Mon, 24 Feb 2020 18:36:23 +0100
Subject: [PATCH] Finished most of the work on the project
---
Cargo.lock | 57 +++++++++++++++++
Cargo.toml | 3 +-
src/ffmpeg_api/api.rs | 21 +++---
src/main.rs | 86 ++++++++++++-------------
src/media_time.rs | 43 +++++++++++++
src/spritesheet.rs | 144 ++++++++++++++++++++++++++++++++++++++++++
src/webvtt.rs | 51 +++++++++++++++
7 files changed, 352 insertions(+), 53 deletions(-)
create mode 100644 src/media_time.rs
create mode 100644 src/spritesheet.rs
create mode 100644 src/webvtt.rs
diff --git a/Cargo.lock b/Cargo.lock
index d251260..510aef4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -583,6 +583,17 @@ dependencies = [
"inflate",
]
+[[package]]
+name = "proc-macro-hack"
+version = "0.5.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ecd45702f76d6d3c75a80564378ae228a85f0b59d2f3ed43c91b4a69eb2ebfc5"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
[[package]]
name = "proc-macro2"
version = "1.0.8"
@@ -676,6 +687,17 @@ dependencies = [
"semver",
]
+[[package]]
+name = "rustversion"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3bba175698996010c4f6dce5e7f173b6eb781fce25d2cfc45e27091ce0b79f6"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
[[package]]
name = "scoped_threadpool"
version = "0.1.9"
@@ -787,6 +809,7 @@ dependencies = [
"fraction",
"image",
"log",
+ "time",
]
[[package]]
@@ -800,6 +823,40 @@ dependencies = [
"miniz_oxide",
]
+[[package]]
+name = "time"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3043ac959c44dccc548a57417876c8fe241502aed69d880efc91166c02717a93"
+dependencies = [
+ "libc",
+ "rustversion",
+ "time-macros",
+ "winapi",
+]
+
+[[package]]
+name = "time-macros"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ae9b6e9f095bc105e183e3cd493d72579be3181ad4004fceb01adbe9eecab2d"
+dependencies = [
+ "proc-macro-hack",
+ "time-macros-impl",
+]
+
+[[package]]
+name = "time-macros-impl"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e987cfe0537f575b5fc99909de6185f6c19c3ad8889e2275e686a873d0869ba1"
+dependencies = [
+ "proc-macro-hack",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
[[package]]
name = "unicode-width"
version = "0.1.7"
diff --git a/Cargo.toml b/Cargo.toml
index 3c685b5..094af42 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,4 +12,5 @@ log = "0.4.8"
image = "0.23.0"
enum_primitive = "0.1.1"
failure = "0.1.6"
-fraction = "0.6.2"
\ No newline at end of file
+fraction = "0.6.2"
+time = "0.2"
\ No newline at end of file
diff --git a/src/ffmpeg_api/api.rs b/src/ffmpeg_api/api.rs
index 217da9b..b1de9f2 100644
--- a/src/ffmpeg_api/api.rs
+++ b/src/ffmpeg_api/api.rs
@@ -5,6 +5,7 @@ use std::marker::PhantomData;
use fraction::Fraction;
use crate::ffmpeg_api::enums::*;
+use crate::media_time;
pub struct AVFormatContext {
base: *mut ffi::AVFormatContext,
@@ -62,7 +63,7 @@ impl<'a> AVFormatContext {
}).find(predicate)
}
- pub fn read_frame(&/*TODO:mut*/ self, packet: &mut AVPacket) -> Result<(), failure::Error> {
+ pub fn read_frame(&self, packet: &mut AVPacket) -> Result<(), failure::Error> {
match unsafe { ffi::av_read_frame(self.base, packet.base) } {
0 => Ok(()),
errno => Err(failure::format_err!("Error while decoding frame: {}", errno))
@@ -292,16 +293,11 @@ impl<'a> AVStream<'a> {
)
}
- pub fn timestamp(self: &AVStream<'a>, timestamp: i64) -> std::time::Duration {
- std::time::Duration::from_millis(
- 1000 *
- timestamp as u64 *
- self.base.time_base.num as u64 /
- self.base.time_base.den as u64
- )
+ pub fn timestamp(self: &AVStream<'a>, timestamp: i64) -> media_time::MediaTime {
+ media_time::MediaTime::from_rational(timestamp, self.time_base())
}
- pub fn duration(&self) -> std::time::Duration {
+ pub fn duration(&self) -> media_time::MediaTime {
self.timestamp(self.base.duration)
}
@@ -324,6 +320,13 @@ impl<'a> AVStream<'a> {
)
}
+ pub fn display_aspect_ratio(&self) -> Fraction {
+ Fraction::new(
+ self.base.display_aspect_ratio.num as u32,
+ self.base.display_aspect_ratio.den as u32,
+ )
+ }
+
pub fn codec_parameters(&self) -> AVCodecParameters {
AVCodecParameters::new(unsafe { self.base.codecpar.as_mut() }.expect("not null"), self)
}
diff --git a/src/main.rs b/src/main.rs
index 42f88c1..dafa6b1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,17 +1,19 @@
+#![allow(dead_code)]
+
pub(crate) mod ffmpeg_api;
+pub(crate) mod webvtt;
+pub(crate) mod media_time;
+pub(crate) mod spritesheet;
use crate::ffmpeg_api::enums::*;
use crate::ffmpeg_api::api::*;
-use image::{ImageBuffer, RgbImage};
+use image::{ImageBuffer};
fn main() -> Result<(), std::io::Error> {
let mut before = std::time::SystemTime::now();
- let output = "/root/spritesheets";
- let input = "/var/lib/data/jellyfin-media/shows/Star Trek: Picard/Season 01/S01E01 - Remembrance.mp4";
-
- //let input = "/home/janne/Workspace/justflix/data/video.mp4";
- //let output = "/home/janne/Workspace/justflix/data/spritesheets";
+ let input = "/home/janne/Workspace/justflix/data/video.mp4";
+ let output = "/home/janne/Workspace/justflix/data/spritesheets";
let mut avformat_context = AVFormatContext::new().unwrap_or_else(|error| {
panic!("Could not allocate a context to process the video: {:?}", error)
@@ -20,9 +22,11 @@ fn main() -> Result<(), std::io::Error> {
panic!("Could not open video input: {:?}", error)
});
- let x = 5;
- let y = 5;
- let mut spritesheet: RgbImage = ImageBuffer::new(160 * x, 90 * x);
+ let mut spritesheet_manager = spritesheet::SpritesheetManager::new(
+ 160,
+ 5, 5,
+ output,
+ );
let mut stream: AVStream = avformat_context.find_stream(|stream| {
stream.codec_parameters().codec_type() == AVMediaType::Video
@@ -45,9 +49,6 @@ fn main() -> Result<(), std::io::Error> {
let mut output_frame = AVFrame::new().unwrap_or_else(|error| {
panic!("Could not create output frame: {:?}", error)
});
- output_frame.init(160, 90, AVPixelFormat::RGB24).unwrap_or_else(|error| {
- panic!("Could not init output frame: {:?}", error)
- });
if codec_parameters.codec_type() == AVMediaType::Video {
let mut codec_context = AVCodecContext::new(&local_codec).unwrap_or_else(|error| {
@@ -68,8 +69,6 @@ fn main() -> Result<(), std::io::Error> {
panic!("Could not create input frame: {:?}", error)
});
- let mut i = 0;
-
println!("Time: {:#?}", before.elapsed().unwrap());
before = std::time::SystemTime::now();
@@ -82,7 +81,7 @@ fn main() -> Result<(), std::io::Error> {
});
while codec_context.out_frame(&mut frame).is_ok() {
println!(
- "Frame {}: {:?} @ {}",
+ "Frame {}: {} @ {}",
frame.coded_picture_number(),
stream.timestamp(frame.pts()),
frame.key_frame()
@@ -90,47 +89,48 @@ fn main() -> Result<(), std::io::Error> {
println!("Reading Time: {:#?}", before.elapsed().unwrap());
before = std::time::SystemTime::now();
- scale_context.reinit(&frame, &output_frame, SwsScaler::FastBilinear).unwrap_or_else(|error| {
- panic!("Could not reinit scale context: {:?}", error)
- });
+ if !spritesheet_manager.initialized() {
+ spritesheet_manager.initialize(frame.width() as u32, frame.height() as u32);
+ output_frame.init(
+ spritesheet_manager.sprite_width() as i32,
+ spritesheet_manager.sprite_height() as i32,
+ AVPixelFormat::RGB24,
+ ).unwrap_or_else(|error| {
+ panic!("Could not init output frame: {:?}", error)
+ });
+ scale_context.reinit(
+ &frame,
+ &output_frame,
+ SwsScaler::FastBilinear,
+ ).unwrap_or_else(|error| {
+ panic!("Could not reinit scale context: {:?}", error)
+ });
+ }
+
scale_context.scale(&frame, &mut output_frame);
println!("Processing Time: {:#?}", before.elapsed().unwrap());
before = std::time::SystemTime::now();
- let current: RgbImage = ImageBuffer::from_raw(160, 90, output_frame.data(0).to_vec()).unwrap();
-
- image::imageops::overlay(
- &mut spritesheet,
- ¤t,
- (i % x) * 160,
- ((i / x) % y) * 90,
+ spritesheet_manager.add_image(
+ stream.timestamp(frame.pts()),
+ ImageBuffer::from_raw(
+ output_frame.width() as u32,
+ output_frame.height() as u32,
+ output_frame.data(0).to_vec(),
+ ).unwrap_or_else(|| {
+ panic!("Could not process frame")
+ }),
);
println!("Writing Time: {:#?}", before.elapsed().unwrap());
before = std::time::SystemTime::now();
-
- i += 1;
-
- if i % (x * y) == 0 {
- spritesheet.save(format!("{}/spritesheet_{}.png", output, (i / (x * y)) - 1)).unwrap_or_else(|error| {
- panic!("Could not write spritesheet: {}", error)
- });
- spritesheet = ImageBuffer::new(160 * x, 90 * x);
- println!("Writing Time: {:#?}", before.elapsed().unwrap());
- before = std::time::SystemTime::now();
- }
}
}
}
- if i % (x * y) != 0 {
- spritesheet.save(format!("{}/spritesheet_{}.png", output, i / (x * y))).unwrap_or_else(|error| {
- panic!("Could not write spritesheet: {}", error)
- });
- println!("Writing Time: {:#?}", before.elapsed().unwrap());
- before = std::time::SystemTime::now();
- }
+ spritesheet_manager.end_frame(stream.duration());
+ spritesheet_manager.save();
}
Ok(())
diff --git a/src/media_time.rs b/src/media_time.rs
new file mode 100644
index 0000000..f5d257c
--- /dev/null
+++ b/src/media_time.rs
@@ -0,0 +1,43 @@
+use fraction::Fraction;
+
+#[derive(Copy,Clone,Debug)]
+pub struct MediaTime(time::Duration);
+
+impl MediaTime {
+ pub fn from_rational(timestamp: i64, base: Fraction) -> MediaTime {
+ let num: u64 = *base.numer().unwrap_or_else(|| {
+ panic!("time base of unusable format")
+ });
+ let den: u64 = *base.denom().unwrap_or_else(|| {
+ panic!("time base of unusable format")
+ });
+
+ MediaTime(time::Duration::milliseconds(
+ 1000 * timestamp * num as i64 / den as i64
+ ))
+ }
+
+ pub fn from_millis(timestamp: i64) -> MediaTime {
+ MediaTime(time::Duration::milliseconds(timestamp))
+ }
+
+ pub fn from_seconds(timestamp: i64) -> MediaTime {
+ MediaTime(time::Duration::seconds(timestamp))
+ }
+}
+
+impl std::fmt::Display for MediaTime {
+ #[inline(always)]
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ let z = self.0.subsec_milliseconds();
+ let s = self.0.whole_seconds() % 60;
+ let m = self.0.whole_seconds() / 60 % 60;
+ let h = self.0.whole_seconds() / 3600;
+
+ if h == 0 {
+ write!(f, "{:02}:{:02}.{:03}", m, s, z)
+ } else {
+ write!(f, "{:02}:{:02}:{:02}.{:03}", h, m, s, z)
+ }
+ }
+}
\ No newline at end of file
diff --git a/src/spritesheet.rs b/src/spritesheet.rs
new file mode 100644
index 0000000..c55a2f5
--- /dev/null
+++ b/src/spritesheet.rs
@@ -0,0 +1,144 @@
+use image::{RgbImage, ImageBuffer};
+use crate::media_time::MediaTime;
+use crate::webvtt::{WebVTTFile, WebVTTCue};
+
+pub struct SpritesheetManager {
+ num_horizontal: u32,
+ num_vertical: u32,
+ max_side: u32,
+ sprite_width: u32,
+ sprite_height: u32,
+ spritesheet: RgbImage,
+ current_image: u32,
+ last_timestamp: MediaTime,
+ metadata: WebVTTFile,
+ output_path: std::string::String,
+ initialized: bool,
+}
+
+impl SpritesheetManager {
+ pub fn new<T: AsRef<str>>(max_side: u32, num_horizontal: u32, num_vertical: u32, output_path: T) -> SpritesheetManager {
+ SpritesheetManager {
+ num_horizontal,
+ num_vertical,
+ max_side,
+ sprite_width: 0,
+ sprite_height: 0,
+ spritesheet: ImageBuffer::new(0, 0),
+ current_image: 0,
+ last_timestamp: MediaTime::from_millis(0),
+ metadata: WebVTTFile::new(),
+ output_path: std::string::String::from(output_path.as_ref()),
+ initialized: false,
+ }
+ }
+
+ pub fn initialize(&mut self, width: u32, height: u32) {
+ if width >= height {
+ self.sprite_width = self.max_side;
+ self.sprite_height = self.sprite_width * height / width;
+ } else {
+ self.sprite_height = self.max_side;
+ self.sprite_width = self.sprite_height * width / height;
+ }
+ self.reinit_buffer();
+ self.initialized = true;
+ }
+
+ fn reinit_buffer(&mut self) {
+ self.spritesheet = ImageBuffer::new(
+ self.sprite_width * self.num_horizontal,
+ self.sprite_height * self.num_vertical,
+ );
+ }
+
+ pub fn initialized(&self) -> bool {
+ self.initialized
+ }
+
+ pub fn sprite_width(&self) -> u32 {
+ self.sprite_width
+ }
+
+ pub fn sprite_height(&self) -> u32 {
+ self.sprite_height
+ }
+
+ fn sprite_index(&self, current: u32) -> u32 {
+ current % (self.num_horizontal * self.num_vertical)
+ }
+
+ fn spritesheet_index(&self, current: u32) -> u32 {
+ current / (self.num_horizontal * self.num_vertical)
+ }
+
+ fn x(&self, current: u32) -> u32 {
+ let index = current % self.num_horizontal;
+ index * self.sprite_width
+ }
+
+ fn y(&self, current: u32) -> u32 {
+ let index = (current / self.num_horizontal) % self.num_vertical;
+ index * self.sprite_height
+ }
+
+ pub fn add_image(&mut self, timestamp: MediaTime, image: RgbImage) {
+ if image.width() != self.sprite_width || image.height() != self.sprite_height {
+ panic!(
+ "Wrong image size: {}x{}, but expected {}x{}",
+ image.width(), image.height(),
+ self.sprite_width, self.sprite_height
+ )
+ }
+
+ let x = self.x(self.current_image);
+ let y = self.y(self.current_image);
+ image::imageops::overlay(
+ &mut self.spritesheet,
+ &image,
+ x, y,
+ );
+
+ if self.current_image != 0 {
+ self.end_frame(timestamp);
+ }
+
+ self.last_timestamp = timestamp;
+ self.current_image += 1;
+
+ if self.sprite_index(self.current_image) == 0 {
+ self.save();
+ }
+ }
+
+ pub fn end_frame(&mut self, timestamp: MediaTime) {
+ self.metadata.add(WebVTTCue::new(
+ self.last_timestamp,
+ timestamp,
+ format!(
+ "spritesheet_{}.jpg#xywh={},{},{},{}",
+ self.spritesheet_index(self.current_image - 1),
+ self.x(self.current_image - 1),
+ self.y(self.current_image - 1),
+ self.sprite_width,
+ self.sprite_height
+ ),
+ ));
+ }
+
+ fn save_spritesheet(&mut self) {
+ self.spritesheet.save(
+ format!("{}/spritesheet_{}.png", self.output_path, self.spritesheet_index(self.current_image))
+ ).unwrap_or_else(|error| {
+ panic!("Could not write spritesheet: {}", error)
+ });
+ self.reinit_buffer();
+ }
+
+ pub fn save(&mut self) {
+ self.save_spritesheet();
+ self.metadata.save(format!("{}/spritesheet.vtt", self.output_path)).unwrap_or_else(|error| {
+ panic!("Could not write spritesheet metadata: {}", error)
+ });
+ }
+}
\ No newline at end of file
diff --git a/src/webvtt.rs b/src/webvtt.rs
new file mode 100644
index 0000000..88deb4d
--- /dev/null
+++ b/src/webvtt.rs
@@ -0,0 +1,51 @@
+use crate::media_time::MediaTime;
+use std::fs::File;
+use std::io::prelude::*;
+use std::io::LineWriter;
+
+pub struct WebVTTFile {
+ cues: Vec<WebVTTCue>
+}
+
+pub struct WebVTTCue {
+ start: MediaTime,
+ end: MediaTime,
+ payload: std::string::String,
+}
+
+impl WebVTTFile {
+ pub fn new() -> WebVTTFile {
+ WebVTTFile {
+ cues: Vec::new()
+ }
+ }
+
+ pub fn add(&mut self, cue: WebVTTCue) {
+ self.cues.push(cue);
+ }
+
+ pub fn save(&self, path: std::string::String) -> Result<(), std::io::Error> {
+ let file = File::create(path)?;
+ let mut file = LineWriter::new(file);
+ file.write_all(b"WEBVTT\n\n")?;
+ for cue in &self.cues {
+ cue.save(&mut file)?;
+ }
+ file.flush()?;
+ Ok(())
+ }
+}
+
+impl WebVTTCue {
+ pub fn new(start: MediaTime, end: MediaTime, payload: std::string::String) -> WebVTTCue {
+ WebVTTCue { start, end, payload }
+ }
+
+ fn save(&self, writer: &mut LineWriter<File>) -> Result<(), std::io::Error>{
+ writer.write_all(format!("{} --> {}\n", self.start, self.end).as_bytes())?;
+ writer.write_all(self.payload.as_bytes())?;
+ writer.write_all(b"\n\n")?;
+
+ Ok(())
+ }
+}
--
GitLab