use super::pds::{Palette, PaletteEntry};
use crate::image::{ImageSize, ToImage, ToOcrImage, ToOcrImageOpt};
use image::{ImageBuffer, Luma, LumaA, Pixel, Primitive};
use std::io::{ErrorKind, Read as _};
type PixelConversion<TargetColor> = fn(&PaletteEntry) -> TargetColor;
#[derive(Clone)]
pub struct RleEncodedImage {
width: u16,
height: u16,
palette: Palette,
raw: Vec<u8>,
}
impl RleEncodedImage {
#[must_use]
pub const fn new(width: u16, height: u16, palette: Palette, raw: Vec<u8>) -> Self {
Self {
width,
height,
palette,
raw,
}
}
pub fn pixels<D: Primitive>(
&self,
convert: PixelConversion<LumaA<D>>,
) -> RlePixelIterator<'_, LumaA<D>> {
RlePixelIterator {
rle_image: self,
raw_data: &self.raw,
current_color: LumaA([D::DEFAULT_MIN_VALUE, D::DEFAULT_MAX_VALUE]),
default_color: LumaA([D::DEFAULT_MAX_VALUE, D::DEFAULT_MIN_VALUE]), nb_remaining_pixels: 0,
convert,
}
}
}
impl ImageSize for RleEncodedImage {
fn width(&self) -> u32 {
u32::from(self.width)
}
fn height(&self) -> u32 {
u32::from(self.height)
}
}
impl<'a> RleEncodedImage {
#[must_use]
pub fn iter(&'a self) -> RlePixelIterator<'a, LumaA<u8>> {
self.into_iter()
}
}
impl<'a> IntoIterator for &'a RleEncodedImage {
type Item = LumaA<u8>;
type IntoIter = RlePixelIterator<'a, LumaA<u8>>;
fn into_iter(self) -> Self::IntoIter {
RlePixelIterator {
rle_image: self,
raw_data: &self.raw,
current_color: LumaA([
<u8 as Primitive>::DEFAULT_MIN_VALUE,
<u8 as Primitive>::DEFAULT_MAX_VALUE,
]), default_color: LumaA([
<u8 as Primitive>::DEFAULT_MAX_VALUE,
<u8 as Primitive>::DEFAULT_MIN_VALUE,
]), nb_remaining_pixels: 0,
convert: pe_to_luma_a,
}
}
}
fn pe_to_luma_a<P: Primitive>(input: &PaletteEntry) -> LumaA<P> {
let luminance = P::from(input.luminance).unwrap();
let alpha = P::from(input.transparency).unwrap();
LumaA([luminance, alpha])
}
pub struct RleToImage<'a, P, C>
where
P: Pixel<Subpixel = u8>,
C: Fn(LumaA<u8>) -> P,
{
rle_image: &'a RleEncodedImage,
conv_fn: C,
}
impl<'a, P, C> RleToImage<'a, P, C>
where
P: Pixel<Subpixel = u8>,
C: Fn(LumaA<u8>) -> P,
{
pub const fn new(rle_image: &'a RleEncodedImage, conv_fn: C) -> Self {
Self { rle_image, conv_fn }
}
}
impl<P, C> ToImage for RleToImage<'_, P, C>
where
P: Pixel<Subpixel = u8>,
C: Fn(LumaA<u8>) -> P,
{
type Pixel = P;
#[profiling::function]
fn to_image(&self) -> ImageBuffer<P, Vec<u8>>
where
P: Pixel<Subpixel = u8>,
{
let width = self.rle_image.width();
let height = self.rle_image.height();
let pixel_iter = self.rle_image.into_iter();
let buf_size = (width * height) as usize * P::CHANNEL_COUNT as usize;
let mut buf = Vec::with_capacity(buf_size);
pixel_iter
.map(|p| (self.conv_fn)(p))
.for_each(|p| buf.extend_from_slice(p.channels()));
ImageBuffer::<P, Vec<u8>>::from_vec(width, height, buf)
.expect("Failed to create image buffer")
}
}
impl<C> ToOcrImage for RleToImage<'_, Luma<u8>, C>
where
C: Fn(LumaA<u8>) -> Luma<u8>,
{
#[profiling::function]
fn image(&self, opt: &ToOcrImageOpt) -> image::GrayImage {
let width = self.rle_image.width();
let height = self.rle_image.height();
let border = opt.border;
let raw_pixels = self.rle_image.into_iter().collect::<Vec<_>>();
ImageBuffer::from_fn(width + border * 2, height + border * 2, |x, y| {
if x < border || x >= width + border || y < border || y >= height + border {
opt.background_color
} else {
let offset = (y - border) * width + (x - border);
let pixel = raw_pixels[offset as usize];
(self.conv_fn)(pixel)
}
})
}
}
pub struct RlePixelIterator<'a, C> {
rle_image: &'a RleEncodedImage,
raw_data: &'a [u8],
current_color: C,
default_color: C,
nb_remaining_pixels: u16,
convert: PixelConversion<C>,
}
impl<Pix, Sub> Iterator for RlePixelIterator<'_, Pix>
where
Sub: Primitive,
Pix: Copy + Pixel<Subpixel = Sub>,
{
type Item = Pix;
fn next(&mut self) -> Option<Self::Item> {
if self.nb_remaining_pixels > 0 {
self.nb_remaining_pixels -= 1;
Some(self.current_color)
} else if let Some((color_id, nb_pixel)) = self.read_next_pixel() {
let color = if let Some(color) = self.rle_image.palette.get(color_id) {
(self.convert)(color)
} else {
self.default_color
};
self.current_color = color;
self.nb_remaining_pixels = nb_pixel - 1;
Some(self.current_color)
} else {
None }
}
fn size_hint(&self) -> (usize, Option<usize>) {
let nb_pixels = (self.rle_image.width() * self.rle_image.height()) as usize;
(nb_pixels, Some(nb_pixels))
}
}
impl<Pix, Sub> ExactSizeIterator for RlePixelIterator<'_, Pix>
where
Sub: Primitive,
Pix: Copy + Pixel<Subpixel = Sub>,
{
}
impl<C> RlePixelIterator<'_, C> {
fn read_next_pixel(&mut self) -> Option<(u8 /*color */, u16 /*nb_pixels*/)> {
const MARKER: u8 = 0;
const COLOR_0: u8 = 0;
loop {
let mut color: [u8; 1] = [0; 1];
let res = self.raw_data.read_exact(&mut color);
if let Err(err) = res {
if err.kind() == ErrorKind::UnexpectedEof {
return None;
}
}
let next = color[0];
if next == MARKER {
let mut l2 = [0; 1];
self.raw_data.read_exact(&mut l2).unwrap();
if l2[0] == MARKER {
} else {
let byte = l2[0];
let nb_pixels = match CountMarker::from(byte) {
CountMarker::Long => {
let mut l3 = [0; 1];
self.raw_data.read_exact(&mut l3).unwrap();
let count_bytes = [byte & 0b0011_1111, l3[0]];
u16::from_be_bytes(count_bytes)
}
CountMarker::Short => {
let count_bits = byte & 0b0011_1111;
u16::from(u8::from_be(count_bits))
}
};
let color_marker = ColorMarker::from(byte);
let color = match color_marker {
ColorMarker::Color0 => COLOR_0,
ColorMarker::ColorN => {
let mut color = [0; 1];
self.raw_data.read_exact(&mut color).unwrap();
color[0]
}
};
return Some((color, nb_pixels));
}
} else {
return Some((next, 1));
}
}
}
}
enum ColorMarker {
Color0,
ColorN,
}
impl From<u8> for ColorMarker {
fn from(value: u8) -> Self {
if (value & 0b1000_0000) > 0 {
Self::ColorN
} else {
Self::Color0
}
}
}
enum CountMarker {
Short,
Long,
}
impl From<u8> for CountMarker {
fn from(value: u8) -> Self {
if (value & 0b0100_0000) > 0 {
Self::Long
} else {
Self::Short
}
}
}