Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Convert video stickers (webm) to webp #32

Merged
merged 8 commits into from
Aug 27, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/rust.yml
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ jobs:
target
key: "${{runner.os}} Rust ${{steps.rust-toolchain.outputs.cachekey}} Lock ${{hashFiles('Cargo.lock')}}"
- uses: msrd0/install-rlottie-action@v1
- run: sudo apt-get update -y && sudo apt-get install -y libavc1394-dev libavdevice-dev
- run: cargo test --workspace --all-features --release -- --include-ignored
env:
RUST_BACKTRACE: 1
Expand Down
61 changes: 59 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions mstickerlib/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ include = ["/src/**/*.rs", "/LICENSE", "/README.md"]
anyhow = "1.0"
async-trait = "0.1"
derive-getters = "0.3.0"
ffmpeg = { package = "ffmpeg-next", version = "6.0" }
flate2 = "1.0"
futures-util = "0.3.25"
generic-array = { version = "0.14" , features = ["serde"] }
Expand All @@ -34,6 +35,7 @@ tempfile = "3.2"
tokio = { version = "1.21", features = ["fs", "parking_lot", "sync"] }
tokio-stream = { version = "0.1", features = ["io-util"], default-features = false }
url = "2.2"
webp-animation = "0.8"

[dev-dependencies]
tokio = { version = "1.21", features = ["macros"] }
Expand Down
60 changes: 47 additions & 13 deletions mstickerlib/src/image.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use crate::{
database,
matrix::{self, Config, Mxc}
matrix::{self, Config, Mxc},
video::webm2webp
};
use anyhow::anyhow;
use flate2::write::GzDecoder;
Expand Down Expand Up @@ -32,6 +33,19 @@ pub struct Image {
pub height: u32
}

fn rayon_run<F, T>(callback: F) -> T
where
F: FnOnce() -> T + Send,
T: Send,
for<'a> &'a mut T: Send
{
let mut result: Option<T> = None;
rayon::scope(|s| {
s.spawn(|_| result = Some(callback()));
});
result.unwrap()
}

impl Image {
pub fn mime_type(&self) -> anyhow::Result<String> {
Ok(format!(
Expand All @@ -56,18 +70,6 @@ impl Image {
if !self.file_name.ends_with(".tgs") {
return Ok(self);
}
fn rayon_run<F, T>(callback: F) -> T
where
F: FnOnce() -> T + Send,
T: Send,
for<'a> &'a mut T: Send
{
let mut result: Option<T> = None;
rayon::scope(|s| {
s.spawn(|_| result = Some(callback()));
});
result.unwrap()
}

tokio::task::spawn_blocking(move || {
rayon_run(move || {
Expand Down Expand Up @@ -102,6 +104,38 @@ impl Image {
.await?
}

pub async fn convert_webm_if_webp(self, animation_format: Option<AnimationFormat>) -> anyhow::Result<Self> {
match animation_format {
Some(AnimationFormat::Webp) => self.convert_webm2webp().await,
_ => Ok(self)
}
}

/// convert `webm` video stickers to webp, ignore other formats
pub async fn convert_webm2webp(mut self) -> anyhow::Result<Self> {
if !self.file_name.ends_with(".webm") {
return Ok(self);
}

tokio::task::spawn_blocking(move || {
rayon_run(move || {
let mut tmp = tempfile::Builder::new().suffix(".webm").tempfile()?;
tmp.write_all(&self.data)?;
tmp.flush()?;

self.file_name.truncate(self.file_name.len() - 1);
self.file_name += "p";
let (webp, width, height) = webm2webp(&tmp.path())?;
self.data = Arc::new(webp.to_vec());
self.width = width;
self.height = height;

Ok(self)
})
})
.await?
}

///upload image to matrix
/// return mxc_url and true if image was uploaded now; false if it was already uploaded before and exist at the database
pub async fn upload<D>(&self, matrix_config: &Config, database: Option<&D>) -> anyhow::Result<(Mxc, bool)>
Expand Down
11 changes: 11 additions & 0 deletions mstickerlib/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ pub mod database;
pub mod image;
pub mod matrix;
pub mod tg;
mod video;

//mod sub_commands;

Expand All @@ -26,6 +27,13 @@ struct Client {
client: UnsafeCell<Option<reqwest::Client>>
}

// XXX Hacky: We abuse the fact that a client will be exactly once either set or
// created, so we can ensure this function will be called exactly once. Also, the
// HTTP client will always be needed before ffmpeg.
fn init() {
ffmpeg::init().expect("Failed to initialise ffmpeg");
}

impl Client {
pub(crate) async fn get(&self) -> &reqwest::Client {
// safety: this method ensures that the client is set from None to Some exactly once, and the
Expand All @@ -44,10 +52,12 @@ impl Client {
let client = unsafe { self.client.get().as_mut().unwrap() };
if client.is_none() {
*client = Some(reqwest::Client::new());
init();
}
client.as_ref().unwrap()
}
}

pub async fn set_client(client: reqwest::Client) {
#[allow(unused_variables)]
let guard = CLIENT.lock.read();
Expand All @@ -56,6 +66,7 @@ pub async fn set_client(client: reqwest::Client) {
panic!("reqwest client was already set")
}
*lib_client = Some(client);
init();
}

pub async fn get_client() -> &'static reqwest::Client {
Expand Down
13 changes: 2 additions & 11 deletions mstickerlib/src/tg/sticker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use crate::{
matrix::{self, sticker_formats::ponies, Mxc},
CLIENT
};
use anyhow::bail;
use derive_getters::Getters;
use log::warn;
use serde::Deserialize;
Expand Down Expand Up @@ -76,6 +75,8 @@ impl PhotoSize {
.download(tg_config)
.await?
.convert_tgs_if_some(advance_config.animation_format)
.await?
.convert_webm_if_webp(advance_config.animation_format)
.await?;
#[cfg(feature = "log")]
info!(" upload sticker{thumb:<10} {pack_name}:{positon:02} {emoji}");
Expand Down Expand Up @@ -129,16 +130,6 @@ impl Sticker {
where
D: crate::database::Database
{
if self.is_video {
#[cfg(feature = "log")]
info!(
" skip Sticker {}:{:02} {}, is a video",
self.pack_name,
self.positon,
self.emoji.as_deref().unwrap_or_default()
);
bail!("sticker is video")
}
#[cfg(feature = "log")]
info!(
"download sticker {}:{:02} {}",
Expand Down
13 changes: 6 additions & 7 deletions mstickerlib/src/tg/stickerpack.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,6 @@ impl StickerPack {
{
#[cfg(feature = "log")]
info!("import Telegram stickerpack {}({})", self.title, self.name);
#[cfg(feature = "log")]
if self.is_video {
warn!(
"sticker pack {} includes video stickers. Import of video stickers is not supported and will be skipped.",
self.name
);
}

let stickers_import_futures = self
.stickers
Expand Down Expand Up @@ -163,4 +156,10 @@ mod tests {
async fn import_none() {
import("NSanimated", None).await;
}

#[tokio::test]
#[ignore]
async fn import_video_pack_webp() {
import("pingu_animated", Some(AnimationFormat::Webp)).await;
}
}
62 changes: 62 additions & 0 deletions mstickerlib/src/video.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
//! This module deals with translating telegram's video stickers to webp animations.

use ffmpeg::{
codec::Context as CodecContext,
decoder,
format::{self, Pixel},
media::Type,
software::scaling::{context::Context as ScalingContext, flag::Flags},
util::frame::video::Video
};
use std::path::Path;
use webp_animation::{Encoder, WebPData};

pub(crate) fn webm2webp<P: AsRef<Path>>(file: &P) -> anyhow::Result<(WebPData, u32, u32)> {
// heavily inspired by
// https://github.com/zmwangx/rust-ffmpeg/blob/master/examples/dump-frames.rs

let mut ictx = format::input(file)?;
let input = ictx.streams().best(Type::Video).ok_or(ffmpeg::Error::StreamNotFound)?;

let video_stream_index = input.index();
let ctx_decoder = CodecContext::from_parameters(input.parameters())?;
let mut decoder = ctx_decoder.decoder().video()?;

let mut scaler = ScalingContext::get(
decoder.format(),
decoder.width(),
decoder.height(),
Pixel::RGBA,
decoder.width(),
decoder.height(),
Flags::BILINEAR
)?;

let mut encoder = Encoder::new((decoder.width(), decoder.height()))?;
let mut timestamp = 0;
let frame_rate = input.rate();
let time_per_frame = frame_rate.1 * 1000 / frame_rate.0;
let mut receive_and_process_decoded_frames = |decoder: &mut decoder::Video| -> anyhow::Result<()> {
let mut decoded = Video::empty();
while decoder.receive_frame(&mut decoded).is_ok() {
let mut rgba_frame = Video::empty();
scaler.run(&decoded, &mut rgba_frame)?;

encoder.add_frame(rgba_frame.data(0), timestamp)?;
timestamp += time_per_frame;
}
Ok(())
};

for (stream, packet) in ictx.packets() {
if stream.index() == video_stream_index {
decoder.send_packet(&packet)?;
receive_and_process_decoded_frames(&mut decoder)?;
}
}
decoder.send_eof()?;
receive_and_process_decoded_frames(&mut decoder)?;

let webp = encoder.finalize(timestamp)?;
Ok((webp, decoder.width(), decoder.height()))
}