tikz-gallery-generator

Custum build of stapix for tikz.pablopie.xyz

Name	Size	Mode
..
src/image.rs	19K	-rw-r--r--

//! Efficient-ish routines for parsing JPEG/PNG and encoding WebP
//!
//! We avoid using the `image` crate because it is bloated, yet we do use some
//! libraries:
//!
//! * `libwebp-sys` (we don't use `webp` because it requires `image`
//! * `zune-jpeg`   (used internaly by `image`)
//! * `zune-png`
//!
//! The `zune-` crates implement SIMD optimizations.
//!
//! Working with these crates directly also allows us the oportunity to
//! streamline pixel format conversions. For example, when parsing JPEG we now
//! feed the YUV pixel data directly to `libwebp` instead of first converting
//! it to  and then converting it back to YUV (which is what `image` would
//! have done).
//!
//! ## Possible improvements
//!
//! * SIMD optimizations           (doesn't seem necessary for now)
//! * use the GPU for downsampling (doesn't seem worth the coplexity/overhead)
#![allow(clippy::identity_op)]
#![allow(clippy::needless_range_loop)]

use std::{
  slice,
  mem::{self, MaybeUninit},
  io::{self, Write, Read, BufReader},
  path::Path,
  fs::File,
};
use zune_core::{
  colorspace::ColorSpace,
  options::DecoderOptions,
  bytestream::ZCursor,
  result::DecodingResult,
};
use zune_jpeg::JpegDecoder;
use zune_png::PngDecoder;
use libwebp_sys::WebPEncodingError;

use crate::create_file;

#[derive(Debug)]
pub struct Image {
  width:     usize,
  height:    usize,

  pixels: PixelBuffer,
}

#[derive(Debug)]
#[allow(clippy::upper_case_acronyms)]
enum PixelBuffer {
  YCbCr {
    y:  Vec<u8>,
    cb: Vec<u8>,
    cr: Vec<u8>,
    a:  Option<Vec<u8>>,
    uv_stride: usize,
  },
  BGRA(Vec<u8>),
}

/// Parses the Y Cb Cr data from a JPEG, downsampling to fit `TARGET_HEIGHT`
/// along the way
pub fn parse_jpeg(path: &Path, target_height: usize) -> Result<Image, ()> {
  let options = DecoderOptions::default()
    .jpeg_set_out_colorspace(ColorSpace::YCbCr);

  let f = BufReader::new(open_file(path).map_err(|_| ())?);
  let mut decoder = JpegDecoder::new_with_options(f, options);

  let pixels = match decoder.decode() {
    Ok(pixels) => pixels,
    Err(e)     => {
      errorln!("Couldn't parse {path:?}: {e}");
      return Err(());
    }
  };

  // ==========================================================================
  let info = decoder.info().expect("should have already parsed JPEG header");
  let options = decoder.options();

  let src_width  = info.width  as usize;
  let src_height = info.height as usize;
  let target_width = src_width * target_height / src_height;

  let (width, height) = if src_height > target_height {
    (target_width, target_height)
  } else {
    (src_width, src_height)
  };

  let uv_width  = width.div_ceil(2);
  let uv_height = height.div_ceil(2);

  assert_expected_pixels_len(src_width, src_height, 3, pixels.len());

  let colorspace = options.jpeg_get_out_colorspace();
  debug_assert!(colorspace == ColorSpace::YCbCr,
                "unexpected colorspace when parsing JPEG: {colorspace:?}");

  // ==========================================================================
  let [y] = unsafe {
    load_and_downsample_channels_linear_checked::<1>(
      &pixels,
      src_width, src_height,
      3,
      target_width, target_height,
    )
  };

  let [cb, cr] = unsafe {
    load_and_downsample_channels_linear::<2>(
      &pixels,
      src_width, src_height,
      3, 1,
      uv_width, uv_height,
    )
  };

  // ==========================================================================
  let y  = unsafe { mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(y)  };
  let cb = unsafe { mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(cb) };
  let cr = unsafe { mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(cr) };

  Ok(Image {
    width,
    height,

    pixels: PixelBuffer::YCbCr { y, cb, cr, a: None, uv_stride: uv_width, },
  })
}

/// Parses the Y Cb Cr data from a PNG, downsampling to fit `TARGET_HEIGHT`
/// along the way
pub fn parse_png(path: &Path, target_height: usize) -> Result<Image, ()> {
  let options = DecoderOptions::default()
    .png_set_strip_to_8bit(true);

  let mut f = open_file(path).map_err(|_| ())?;
  let size: usize = f.metadata()
    .expect("OS should support file metadata")
    .len() as _;
  let mut f_contents = Vec::with_capacity(size);
  f.read_to_end(&mut f_contents)
    .map_err(|e| errorln!("Could not read {path:?}: {e}"))?;
  let f_contents = ZCursor::new(f_contents);

  let mut decoder = PngDecoder::new_with_options(f_contents, options);

  let pixels = match decoder.decode() {
    Ok(DecodingResult::U8(pixels)) => pixels,
    Ok(DecodingResult::U16(_)) => {
      unreachable!("PNG strip was set to 8 bits")
    }
    Ok(_) => {
      unreachable!("PngDecoder::decode should always return DecodingResult::U8 or DecodingResult::U16")
    }
    Err(e) => {
      errorln!("Couldn't parse {path:?}: {e}");
      return Err(());
    }
  };

  // ==========================================================================
  let (src_width, src_height) = decoder
    .dimensions()
    .expect("should already have decoded PNG headers");
  let target_width = src_width * target_height / src_height;

  let (width, height) = if src_height >= target_height {
    (target_width, target_height)
  } else {
    (src_width, src_height)
  };

  // ==========================================================================
  let colorspace = decoder
    .colorspace()
    .expect("should already have decoded PNG headers");

  debug_assert!(
    matches!(
      colorspace,
      ColorSpace::RGB | ColorSpace::RGBA | ColorSpace::Luma | ColorSpace::LumaA
    ),
    "unexpected colorspace when parsing PNG: {colorspace:?}",
  );

  let has_alpha = matches!(colorspace, ColorSpace::RGBA | ColorSpace::LumaA);
  let is_grayscale = matches!(
    colorspace,
    ColorSpace::Luma | ColorSpace::LumaA,
  );

  let mut bps = 1;
  if !is_grayscale { bps += 2; }
  if has_alpha     { bps += 1; }
  assert_expected_pixels_len(src_width, src_height, bps, pixels.len());

  // ==========================================================================
  // handle RGB/RGBA input

  if !is_grayscale {
    let bgra_data = if src_height < target_height {
      unsafe { load_bgra_from_rgba(&pixels, src_width, src_height, has_alpha) }
    } else {
      unsafe {
        load_and_downsample_bgra_from_rgba(
          &pixels,
          src_width,    src_height,
          has_alpha,
          target_width, target_height,
        )
      }
    };

    return Ok(Image { width, height, pixels: PixelBuffer::BGRA(bgra_data), });
  }

  // ==========================================================================
  // handle grayscale input

  let uv_width  = width.div_ceil(2);
  let uv_height = height.div_ceil(2);
  let uv_stride = uv_width;
  let uv_len = uv_width * uv_height;

  let cb = vec![128; uv_len];
  let cr = vec![128; uv_len];

  if !has_alpha {
    if src_height < target_height {
      return Ok(Image {
        width,
        height,

        pixels: PixelBuffer::YCbCr { y: pixels, cb, cr, a: None, uv_stride, }
      });
    }

    let [y] = unsafe {
      load_and_downsample_channels_linear::<1>(
        &pixels,
        src_width,    src_height,
        1, 0,
        target_width, target_height,
      )
    };
    let y = unsafe { mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(y) };

    return Ok(Image {
      width,
      height,

      pixels: PixelBuffer::YCbCr {y, a: None, cb, cr, uv_stride, }
    });
  }

  let [y, a] = unsafe {
    load_and_downsample_channels_linear_checked::<2>(
      &pixels,
      src_width,    src_height,
      2,
      target_width, target_height,
    )
  };
  let y = unsafe { mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(y) };
  let a = unsafe { mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(a) };

  Ok(Image {
    width,
    height,

    pixels: PixelBuffer::YCbCr { y, a: Some(a), cb, cr, uv_stride, },
  })
}

pub fn encode_webp(img: &Image, output_path: &Path) -> Result<(), ()> {
  use libwebp_sys::*;
  const WEBP_IMAGE_QUALITY: f32 = 90.0;

  unsafe {
    // ======================================================================
    let mut config = MaybeUninit::uninit();
    let init = WebPInitConfig(config.as_mut_ptr());
    let mut config: WebPConfig = config.assume_init();
    debug_assert!(init, "libwebp version mismatch");

    config.lossless          = 0; // using lossy compression
    config.alpha_compression = 1; // using lossy compression
    config.quality           = WEBP_IMAGE_QUALITY;

    debug_assert!(WebPValidateConfig(&config) != 0,
                  "WebP config should be valid");

    // ======================================================================
    // NOTE: it makes no sense to re-use ww because WebPMemoryWriterClear(ww)
    //       literally frees ww.mem and sets it NULL
    let mut ww = MaybeUninit::uninit();
    WebPMemoryWriterInit(ww.as_mut_ptr());
    let mut ww: WebPMemoryWriter = ww.assume_init();

    // NOTE: it makes no sence to re-use picture since there is no allocated
    //       data other then its fields
    let mut picture = MaybeUninit::uninit();
    let init = WebPPictureInit(picture.as_mut_ptr());
    let mut picture: WebPPicture = picture.assume_init();
    debug_assert!(init, "libwebp version mismatch");


    picture.width  = img.width  as i32;
    picture.height = img.height as i32;

    // NOTE: setting the planes manually for efficciency
    match &img.pixels {
      PixelBuffer::YCbCr { y, cb, cr, a, uv_stride, } => {
        // picture.colorspace = WebPEncCSP::WEBP_YUV420; // already default
        // picture.use_argb = 0;                         // already default
        picture.y_stride  = img.width  as i32;
        picture.uv_stride = *uv_stride as i32;
        picture.y =  y.as_ptr() as *mut _;
        picture.u = cb.as_ptr() as *mut _;
        picture.v = cr.as_ptr() as *mut _;

        if let Some(a) = a {
          picture.colorspace = WebPEncCSP::WEBP_YUV420A;
          picture.a_stride   = img.width as i32;
          picture.a          = a.as_ptr() as *mut _;
        }
      }
      PixelBuffer::BGRA(bgra_data) => {
        picture.use_argb    = 1;
        picture.argb_stride = img.width as i32;
        picture.argb        = bgra_data.as_ptr() as *mut _;
      }
    }

    debug_assert!(WebPPictureIsView(&picture) != 0,
                  "picture should have been configured as \"view\" by not calling WebPPictureAlloc or any of the \"Import\" routines");

    picture.writer     = Some(WebPMemoryWrite);
    picture.custom_ptr = &mut ww as *mut WebPMemoryWriter as _;

    // ======================================================================
    let status = WebPEncode(&config, &mut picture) != 0;
    let bytes  = if status {
      slice::from_raw_parts(ww.mem, ww.size)
    } else {
      log_webp_memmory_error(picture.error_code, output_path);
      return Err(());
    };

    // ======================================================================
    let mut thumb_file = create_file(output_path).map_err(|_| ())?;
    if let Err(e) = thumb_file.write_all(bytes) {
      errorln!("Couldn't write thumnail to file {output_path:?}: {e}");
      return Err(());
    }

    WebPPictureFree(&mut picture);
    WebPMemoryWriterClear(&mut ww);
  }

  Ok(())
}

#[inline]
unsafe fn load_and_downsample_channels_linear<const CHANNELS: usize>(
  pixels: &[u8],
  src_width: usize, src_height: usize,
  bps: usize, first_channel_offset: usize,
  target_width: usize, target_height: usize,
) -> [Vec<MaybeUninit<u8>>; CHANNELS] {
  debug_assert!(CHANNELS > 0);

  let mut result: [Vec<MaybeUninit<u8>>; CHANNELS] = {
    [const { Vec::new() }; CHANNELS]
  };

  for j in 0..CHANNELS {
    result[j].resize(target_width * target_height, MaybeUninit::uninit());
  }

  let mut dst_i = 0;
  for dst_y in 0..target_height {
    let bottom = (      dst_y * src_height).div_ceil(target_height);
    let top    = ((dst_y + 1) * src_height).div_ceil(target_height);

    for dst_x in 0..target_width {
      let left  = (      dst_x * src_width).div_ceil(target_width);
      let right = ((dst_x + 1) * src_width).div_ceil(target_width);

      // TODO: handle this
      assert!(top != bottom && left != right);
      let area: u16 = ((right - left)*(top - bottom)) as u16;

      let mut acc: [u16; CHANNELS] = [0; CHANNELS];
      let mut i = bottom * src_width + left;
      for _ in bottom..top {
        for _ in left..right {
          for j in 0..CHANNELS {
            acc[j] += pixels[i*bps+first_channel_offset+j] as u16;
          }
          i += 1;
        }
        i += src_width + left;
        i -= right;
      }

      for j in 0..CHANNELS {
        result[j][dst_i] = MaybeUninit::new((acc[j]/area) as u8);
      }

      dst_i += 1;
    }
  }

  result
}

#[inline]
unsafe fn load_channels<const CHANNELS: usize>(
  pixels: &[u8],
  src_width: usize, src_height: usize,
  bps: usize,
) -> [Vec<MaybeUninit<u8>>; CHANNELS] {
  debug_assert!(CHANNELS > 0);

  let mut result = [const { Vec::new() }; CHANNELS];

  for j in 0..CHANNELS {
    result[j].resize(src_width * src_height, MaybeUninit::uninit());
  }

  for j in 0..CHANNELS {
    for i in 0..src_width*src_height {
      result[j][i] = MaybeUninit::new(pixels[i*bps+j]);
    }
  }

  result
}

#[inline]
unsafe fn load_and_downsample_channels_linear_checked<const CHANNELS: usize>(
  pixels: &[u8],
  src_width: usize, src_height: usize,
  bps: usize,
  target_width: usize, target_height: usize,
) -> [Vec<MaybeUninit<u8>>; CHANNELS] {
  if src_height < target_height {
    load_channels::<CHANNELS>(pixels, src_width, src_height, bps)
  } else {
    load_and_downsample_channels_linear::<CHANNELS>(
      pixels,
      src_width, src_height,
      bps, 0,
      target_width, target_height,
    )
  }
}

#[inline]
unsafe fn load_and_downsample_bgra_from_rgba(
  pixels: &[u8],
  src_width: usize, src_height: usize,
  has_alpha: bool,
  target_width: usize, target_height: usize,
) -> Vec<u8> {
  if has_alpha {
    return load_and_downsample_bgra_from_rgba_internal::<true>(
      pixels,
      src_width,    src_height,
      target_width, target_height,
    );
  } else {
    return load_and_downsample_bgra_from_rgba_internal::<false>(
      pixels,
      src_width,    src_height,
      target_width, target_height,
    );
  }

  #[inline]
  unsafe fn load_and_downsample_bgra_from_rgba_internal<const HAS_ALPHA: bool>(
    pixels: &[u8],
    src_width:    usize, src_height:    usize,
    target_width: usize, target_height: usize,
  ) -> Vec<u8> {
    let bps = if HAS_ALPHA { 4 } else { 3 };

    let len = target_width * target_height;
    let mut result = vec![MaybeUninit::uninit(); len * 4];

    let mut dst_i = 0;
    for dst_y in 0..target_height {
      let bottom = (      dst_y * src_height).div_ceil(target_height);
      let top    = ((dst_y + 1) * src_height).div_ceil(target_height);

      for dst_x in 0..target_width {
        let left  = (      dst_x * src_width).div_ceil(target_width);
        let right = ((dst_x + 1) * src_width).div_ceil(target_width);

        // TODO: handle this
        assert!(top != bottom && left != right);

        let area: u16 = ((right - left)*(top - bottom)) as u16;

        #[allow(unused_variables)]
        let mut acc_a: u16 = 0;
        let mut acc_r: u16 = 0;
        let mut acc_g: u16 = 0;
        let mut acc_b: u16 = 0;

        let mut i = bottom * src_width + left;
        for _ in bottom..top {
          for _ in left..right {
            acc_r += pixels[i*bps+0] as u16;
            acc_g += pixels[i*bps+1] as u16;
            acc_b += pixels[i*bps+2] as u16;
            // NOTE: compile time branches should never branch at runtime
            #[allow(unused_assignments)]
            if HAS_ALPHA { acc_a += pixels[i*bps+3] as u16; }
            i += 1;
          }
          i += src_width + left;
          i -= right;
        }

        result[dst_i*4+0] = MaybeUninit::new((acc_b/area) as u8);
        result[dst_i*4+1] = MaybeUninit::new((acc_g/area) as u8);
        result[dst_i*4+2] = MaybeUninit::new((acc_r/area) as u8);

        // NOTE: compile time branches should never branch at runtime
        if HAS_ALPHA {
          result[dst_i*4+3] = MaybeUninit::new((acc_a/area) as u8);
        } else {
          result[dst_i*4+3] = MaybeUninit::new(0xff);
        }

        dst_i += 1;
      }
    }

    mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(result)
  }
}

#[inline]
unsafe fn load_bgra_from_rgba(
  pixels: &[u8],
  src_width: usize, src_height: usize,
  has_alpha: bool,
) -> Vec<u8> {
  if has_alpha {
    return load_bgra_from_rgba_internal::<true>(
      pixels, src_width, src_height,
    );
  } else {
    return load_bgra_from_rgba_internal::<false>(
      pixels, src_width, src_height,
    );
  }

  #[inline]
  unsafe fn load_bgra_from_rgba_internal<const HAS_ALPHA: bool>(
    pixels: &[u8],
    src_width: usize, src_height: usize,
  ) -> Vec<u8> {
    let bps = if HAS_ALPHA { 4 } else { 3 };

    let len = src_width * src_height;
    let mut result = vec![MaybeUninit::uninit(); len * 4];

    for i in 0..len {
      let r = pixels[i*bps+0];
      let g = pixels[i*bps+1];
      let b = pixels[i*bps+2];

      // NOTE: compile time branches should not actually branch at runtime
      let mut a = 0xff;
      if HAS_ALPHA { a = pixels[i*bps+3]; }

      result[i*4+0] = MaybeUninit::new(b);
      result[i*4+1] = MaybeUninit::new(g);
      result[i*4+2] = MaybeUninit::new(r);
      result[i*4+3] = MaybeUninit::new(a);
    }

    mem::transmute::<Vec<MaybeUninit<u8>>, Vec<u8>>(result)
  }
}

#[inline]
fn assert_expected_pixels_len(
  width: usize,
  height: usize,
  bpp: usize,
  pixels_len: usize,
) {
  let expected_pixels_len = width * height * bpp;
  debug_assert!(
    pixels_len == expected_pixels_len,
    "unexpected pixel buffer length: expected {width} × {height} × {bpp} = {expected_pixels_len} but got {pixels_len}",
  );
}

#[inline]
fn log_webp_memmory_error(e: WebPEncodingError, output_path: &Path) {
  match e {
    WebPEncodingError::VP8_ENC_ERROR_OUT_OF_MEMORY => {
      errorln!("Could not generate {output_path:?}: out of memmory");
    }
    WebPEncodingError::VP8_ENC_ERROR_BITSTREAM_OUT_OF_MEMORY => {
      errorln!("Could not generate {output_path:?}: memory error while flushing bits");
    }
    WebPEncodingError::VP8_ENC_ERROR_PARTITION0_OVERFLOW => {
      errorln!("Could not generate {output_path:?}: partition is bigger than 512k");
    }
    WebPEncodingError::VP8_ENC_ERROR_PARTITION_OVERFLOW => {
      errorln!("Could not generate {output_path:?}: partition is bigger than 16M");
    }
    WebPEncodingError::VP8_ENC_ERROR_BAD_WRITE => {
      errorln!("Could not generate {output_path:?}: error while flushing bytes");
    }
    WebPEncodingError::VP8_ENC_ERROR_FILE_TOO_BIG => {
      errorln!("Could not generate {output_path:?}: file is bigger than 4G");
    }
    WebPEncodingError::VP8_ENC_ERROR_USER_ABORT => {
      errorln!("Could not generate {output_path:?}: abort requested by the user");
    }
    e => unreachable!("unexpected WebPEncodingError variant: {e:?}"),
  }
}

fn open_file(path: &Path) -> io::Result<File> {
  File::open(path)
    .map_err(|e| { errorln!("Could not parse {path:?}: {e}"); e })
}