Skip to content

Commit

Permalink
Parallel blur and quantization
Browse files Browse the repository at this point in the history
  • Loading branch information
@kornelski
kornelski committed Dec 4, 2022
1 parent d115651 commit 5fdaca2
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 74 deletions.
6 changes: 3 additions & 3 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ repository = "https://github.com/ImageOptim/gifski"
version = "1.8.1"
autobins = false
edition = "2021"
rust-version = "1.60"
rust-version = "1.63"

[[bin]]
doctest = false
Expand Down Expand Up @@ -86,5 +86,5 @@ strip = true
[package.metadata.docs.rs]
targets = ["x86_64-unknown-linux-gnu"]

# [patch.crates-io]
# ffmpeg-sys-next = { rev = "78458039d5fac99354f9cb078869f3be3f3af5fb", git = "https://github.com/kornelski/rust-ffmpeg-sys-1"}
[patch.crates-io]
ffmpeg-sys-next = { rev = "78458039d5fac99354f9cb078869f3be3f3af5fb", git = "https://github.com/kornelski/rust-ffmpeg-sys-1"}
73 changes: 40 additions & 33 deletions src/denoise.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,19 +131,23 @@ impl<T> Denoiser<T> {
}
}

pub fn push_frame(&mut self, frame: ImgRef<RGBA8>, frame_metadata: T) -> Result<(), WrongSizeError> {
#[cfg(test)]
fn push_frame_test(&mut self, frame: ImgRef<RGBA8>, frame_metadata: T) -> Result<(), WrongSizeError> {
let frame_blurred = smart_blur(frame);
self.push_frame(frame, frame_blurred.as_ref(), frame_metadata)
}

pub fn push_frame(&mut self, frame: ImgRef<RGBA8>, frame_blurred: ImgRef<RGB8>, frame_metadata: T) -> Result<(), WrongSizeError> {
if frame.width() != self.splat.width() || frame.height() != self.splat.height() {
return Err(WrongSizeError);
}

self.metadatas.insert(0, frame_metadata);

let frame_blurred = smart_blur(frame);

self.frames += 1;
// Can't output anything yet
if self.frames < LOOKAHEAD {
self.quick_append(frame, frame_blurred.as_ref());
self.quick_append(frame, frame_blurred);
return Ok(());
}

Expand Down Expand Up @@ -275,7 +279,7 @@ macro_rules! median_channel {
}
}

fn smart_blur(frame: ImgRef<RGBA8>) -> ImgVec<RGB8> {
pub(crate) fn smart_blur(frame: ImgRef<RGBA8>) -> ImgVec<RGB8> {
let mut out = Vec::with_capacity(frame.width() * frame.height());
loop9_img(frame, |_,_, top, mid, bot| {
out.push_in_cap(if mid.curr.a > 0 {
Expand Down Expand Up @@ -364,7 +368,10 @@ fn px<T>(f: Denoised<T>) -> (RGBA8, T) {
fn one() {
let mut d = Denoiser::new(1,1, 100);
let w = RGBA8::new(255,255,255,255);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
let frame = ImgVec::new(vec![w], 1, 1);
let frame_blurred = smart_blur(frame.as_ref());

d.push_frame(frame.as_ref(), frame_blurred.as_ref(), 0).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.flush();
assert_eq!(px(d.pop()), (w, 0));
Expand All @@ -376,8 +383,8 @@ fn two() {
let mut d = Denoiser::new(1,1, 100);
let w = RGBA8::new(254,253,252,255);
let b = RGBA8::new(8,7,0,255);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 1).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 1).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.flush();
assert_eq!(px(d.pop()), (w, 0));
Expand All @@ -390,9 +397,9 @@ fn three() {
let mut d = Denoiser::new(1,1, 100);
let w = RGBA8::new(254,253,252,255);
let b = RGBA8::new(8,7,0,255);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 1).unwrap();
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 1).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.flush();
assert_eq!(px(d.pop()), (w, 0));
Expand All @@ -408,10 +415,10 @@ fn four() {
let w = RGBA8::new(254,253,252,255);
let b = RGBA8::new(8,7,0,255);
let t = RGBA8::new(0,0,0,0);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame(ImgVec::new(vec![t], 1, 1).as_ref(), 1).unwrap();
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 3).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame_test(ImgVec::new(vec![t], 1, 1).as_ref(), 1).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 3).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.flush();
assert_eq!(px(d.pop()), (w, 0));
Expand All @@ -427,12 +434,12 @@ fn five() {
let w = RGBA8::new(254,253,252,255);
let b = RGBA8::new(8,7,0,255);
let t = RGBA8::new(0,0,0,0);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame(ImgVec::new(vec![t], 1, 1).as_ref(), 1).unwrap();
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 3).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame_test(ImgVec::new(vec![t], 1, 1).as_ref(), 1).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 3).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 4).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 4).unwrap();
assert_eq!(px(d.pop()), (w, 0));
d.flush();
assert_eq!(px(d.pop()), (t, 1));
Expand All @@ -449,17 +456,17 @@ fn six() {
let b = RGBA8::new(8,7,0,255);
let t = RGBA8::new(0,0,0,0);
let x = RGBA8::new(4,5,6,255);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 0).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 1).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 1).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), 2).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![t], 1, 1).as_ref(), 3).unwrap();
d.push_frame_test(ImgVec::new(vec![t], 1, 1).as_ref(), 3).unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), 4).unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), 4).unwrap();
assert_eq!(px(d.pop()), (w, 0));
d.push_frame(ImgVec::new(vec![x], 1, 1).as_ref(), 5).unwrap();
d.push_frame_test(ImgVec::new(vec![x], 1, 1).as_ref(), 5).unwrap();
d.flush();
assert_eq!(px(d.pop()), (b, 1));
assert_eq!(px(d.pop()), (b, 2));
Expand All @@ -476,19 +483,19 @@ fn many() {
let w = RGBA8::new(255,254,253,255);
let b = RGBA8::new(1,2,3,255);
let t = RGBA8::new(0,0,0,0);
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), "w0").unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), "w0").unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![w], 1, 1).as_ref(), "w1").unwrap();
d.push_frame_test(ImgVec::new(vec![w], 1, 1).as_ref(), "w1").unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), "b2").unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), "b2").unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), "b3").unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), "b3").unwrap();
assert!(matches!(d.pop(), Denoised::NotYet));
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), "b4").unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), "b4").unwrap();
assert_eq!(px(d.pop()), (w, "w0"));
d.push_frame(ImgVec::new(vec![t], 1, 1).as_ref(), "t5").unwrap();
d.push_frame_test(ImgVec::new(vec![t], 1, 1).as_ref(), "t5").unwrap();
assert_eq!(px(d.pop()), (w, "w1"));
d.push_frame(ImgVec::new(vec![b], 1, 1).as_ref(), "b6").unwrap();
d.push_frame_test(ImgVec::new(vec![b], 1, 1).as_ref(), "b6").unwrap();
assert_eq!(px(d.pop()), (b, "b2"));
d.flush();
assert_eq!(px(d.pop()), (b, "b3"));
Expand Down
85 changes: 47 additions & 38 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ use std::borrow::Cow;
struct InputFrame {
/// The pixels to encode
frame: ImgVec<RGBA8>,
/// The same as above, but with smart blur applied (for denoiser)
frame_blurred: ImgVec<RGB8>,
/// Time in seconds when to display the frame. First frame should start at 0.
presentation_timestamp: f64,
}
Expand Down Expand Up @@ -223,8 +225,10 @@ impl Collector {

pub(crate) fn add_frame_rgba_cow(&self, frame_index: usize, image: Img<Cow<'_, [RGBA8]>>, presentation_timestamp: f64) -> CatResult<()> {
debug_assert!(frame_index == 0 || presentation_timestamp > 0.);
let frame = Self::resized_binary_alpha(image, self.width, self.height)?;
self.queue.push(frame_index, Ok(InputFrame {
frame: Self::resized_binary_alpha(image, self.width, self.height)?,
frame_blurred: smart_blur(frame.as_ref()),
frame,
presentation_timestamp,
}))
}
Expand All @@ -244,8 +248,10 @@ impl Collector {
.map_err(|err| Error::PNG(format!("Can't load {}: {err}", path.display())))?;

let image = Img::new(image.buffer.into(), image.width, image.height);
let frame = Self::resized_binary_alpha(image, width, height)?;
self.queue.push(frame_index, Ok(InputFrame {
frame: Self::resized_binary_alpha(image, width, height)?,
frame_blurred: smart_blur(frame.as_ref()),
frame,
presentation_timestamp}
))
}
Expand Down Expand Up @@ -483,15 +489,15 @@ impl Writer {

let settings_ext = self.settings;
let settings = self.settings.s;
let (quant_queue, quant_queue_recv) = crossbeam_channel::bounded(3);
let (quant_queue, quant_queue_recv) = crossbeam_channel::bounded(2);
let diff_thread = thread::Builder::new().name("diff".into()).spawn(move || {
Self::make_diffs(decode_queue_recv, quant_queue, &settings_ext)
})?;
let (remap_queue, remap_queue_recv) = crossbeam_channel::bounded(8);
let (remap_queue, remap_queue_recv) = ordqueue::new(2);
let quant_thread = thread::Builder::new().name("quant".into()).spawn(move || {
Self::quantize_frames(quant_queue_recv, remap_queue, &settings_ext)
})?;
let (write_queue, write_queue_recv) = crossbeam_channel::bounded(4);
let (write_queue, write_queue_recv) = crossbeam_channel::bounded(2);
let remap_thread = thread::Builder::new().name("remap".into()).spawn(move || {
Self::remap_frames(remap_queue_recv, write_queue, &settings)
})?;
Expand Down Expand Up @@ -531,7 +537,7 @@ impl Writer {
let curr_frame = next_frame.take();
next_frame = inputs.next().transpose()?;

if let Some(InputFrame { frame: image, presentation_timestamp: raw_pts }) = curr_frame {
if let Some(InputFrame { frame, frame_blurred, presentation_timestamp: raw_pts }) = curr_frame {
ordinal_frame_number += 1;

let pts = raw_pts - last_frame_duration.shift_every_pts_by();
Expand All @@ -540,8 +546,8 @@ impl Writer {
}
last_frame_pts = pts;

denoiser.push_frame(image.as_ref(), (ordinal_frame_number, pts, last_frame_duration)).map_err(|_| {
Error::WrongSize(format!("Frame {ordinal_frame_number} has wrong size ({}×{})", image.width(), image.height()))
denoiser.push_frame(frame.as_ref(), frame_blurred.as_ref(), (ordinal_frame_number, pts, last_frame_duration)).map_err(|_| {
Error::WrongSize(format!("Frame {ordinal_frame_number} has wrong size ({}×{})", frame.width(), frame.height()))
})?;
if next_frame.is_none() {
denoiser.flush();
Expand Down Expand Up @@ -572,7 +578,8 @@ impl Writer {
Ok(())
}

fn quantize_frames(inputs: Receiver<DiffMessage>, remap_queue: Sender<RemapMessage>, settings: &SettingsExt) -> CatResult<()> {
fn quantize_frames(inputs: Receiver<DiffMessage>, remap_queue: OrdQueue<RemapMessage>, settings: &SettingsExt) -> CatResult<()> {
minipool::new(3, "quant", move |to_remap| {
let mut inputs = inputs.into_iter().peekable();

let DiffMessage {image: first_frame, ..} = inputs.peek().ok_or(Error::NoFrames)?;
Expand Down Expand Up @@ -609,7 +616,7 @@ impl Writer {
gif::DisposalMethod::Keep
};

let mut importance_map = importance_map.take().unwrap(); // always set at the beginning
let importance_map = importance_map.take().unwrap(); // always set at the beginning

if !prev_frame_keeps || importance_map.iter().any(|&px| px > 0) {
if prev_frame_keeps {
Expand All @@ -621,36 +628,40 @@ impl Writer {
}
}

let needs_transparency = consecutive_frame_num > 0 || (consecutive_frame_num == 0 && first_frame_has_transparency);
let (liq, remap, liq_image) = Self::quantize(image, &importance_map, consecutive_frame_num == 0, needs_transparency, prev_frame_keeps, settings)?;
let max_loss = settings.gifsicle_loss();
for imp in &mut importance_map {
// encoding assumes rgba background looks like encoded background, which is not true for lossy
*imp = ((256 - u32::from(*imp)) * max_loss / 256).min(255) as u8;
}

let end_pts = if let Some(&DiffMessage { pts: next_pts, .. }) = inputs.peek() {
next_pts
} else {
pts + frame_duration
};
debug_assert!(end_pts > 0.);

remap_queue.send(RemapMessage {
ordinal_frame_number,
end_pts,
dispose,
liq, remap,
liq_image,
})?;
to_remap.send((end_pts, image, importance_map, ordinal_frame_number, consecutive_frame_num, dispose, first_frame_has_transparency, prev_frame_keeps))?;

consecutive_frame_num += 1;
prev_frame_keeps = dispose == gif::DisposalMethod::Keep;
}
}
Ok(())
}, move |(end_pts, image, mut importance_map, ordinal_frame_number, consecutive_frame_num, dispose, first_frame_has_transparency, prev_frame_keeps)| {
let needs_transparency = consecutive_frame_num > 0 || (consecutive_frame_num == 0 && first_frame_has_transparency);
let (liq, remap, liq_image) = Self::quantize(image, &importance_map, consecutive_frame_num == 0, needs_transparency, prev_frame_keeps, settings).unwrap();
let max_loss = settings.gifsicle_loss();
for imp in &mut importance_map {
// encoding assumes rgba background looks like encoded background, which is not true for lossy
*imp = ((256 - u32::from(*imp)) * max_loss / 256).min(255) as u8;
}

remap_queue.push(consecutive_frame_num as usize, RemapMessage {
ordinal_frame_number,
end_pts,
dispose,
liq, remap,
liq_image,
})
})
}

fn remap_frames(inputs: Receiver<RemapMessage>, write_queue: Sender<FrameMessage>, settings: &Settings) -> CatResult<()> {
fn remap_frames(inputs: OrdQueueIter<RemapMessage>, write_queue: Sender<FrameMessage>, settings: &Settings) -> CatResult<()> {
let mut inputs = inputs.into_iter().peekable();
let first_frame = inputs.peek().ok_or(Error::NoFrames)?;
let mut screen = gif_dispose::Screen::new(first_frame.liq_image.width(), first_frame.liq_image.height(), RGBA8::new(0, 0, 0, 0), None);
Expand Down Expand Up @@ -681,21 +692,19 @@ impl Writer {

screen_after_dispose.then_blit(Some(&image8_pal), dispose, left, top as _, image8.as_ref(), transparent_index)?;

let frame = GIFFrame {
left,
top,
screen_width,
screen_height,
image: image8,
pal: image8_pal,
transparent_index,
dispose,
};

write_queue.send(FrameMessage {
ordinal_frame_number,
end_pts,
frame,
frame: GIFFrame {
left,
top,
screen_width,
screen_height,
image: image8,
pal: image8_pal,
transparent_index,
dispose,
},
})?;

is_first_frame = false;
Expand Down

0 comments on commit 5fdaca2

Please sign in to comment.