use chromatic::Colour;
use ndarray::Array2;
use num_traits::Float;
use png::{BitDepth, ColorType, Decoder, Encoder};
use std::{
fs::File,
io::{BufReader, BufWriter, Read, Write},
path::Path,
};
use crate::{Image, PngError};
impl<C, T, const N: usize> Image<C, T, N> for Array2<C>
where
C: Colour<T, N> + Copy,
T: Float + Send + Sync,
{
fn load<P: AsRef<Path>>(path: P) -> Result<Self, PngError> {
let rd = BufReader::new(File::open(path)?);
Self::read(rd)
}
fn read<R: Read>(reader: R) -> Result<Self, PngError> {
let mut reader = Decoder::new(reader).read_info()?;
let info = reader.info();
let (w, h) = (info.width as usize, info.height as usize);
if info.bit_depth != BitDepth::Eight {
return Err(PngError::UnsupportedBitDepth(info.bit_depth));
}
let expected = match N {
1 => ColorType::Grayscale,
2 => ColorType::GrayscaleAlpha,
3 => ColorType::Rgb,
4 => ColorType::Rgba,
_ => return Err(PngError::InvalidChannelCount),
};
if !match_colour_types(info.color_type, expected) {
return Err(PngError::UnsupportedColourType(info.color_type));
}
let mut buf = vec![0; reader.output_buffer_size()];
reader.next_frame(&mut buf)?;
let pixels = buf
.chunks_exact(N)
.map(|chunk| {
let mut arr = [0u8; N];
arr.copy_from_slice(chunk);
C::from_bytes(arr)
})
.collect::<Vec<_>>();
Array2::from_shape_vec((h, w), pixels).map_err(|_| PngError::InvalidData)
}
fn save<P: AsRef<Path>>(&self, path: P) -> Result<(), PngError> {
let wr = BufWriter::new(File::create(path)?);
Self::write(self, wr)
}
fn write<W: Write>(&self, mut writer: W) -> Result<(), PngError> {
let (h, w) = self.dim();
let colour = match N {
1 => ColorType::Grayscale,
2 => ColorType::GrayscaleAlpha,
3 => ColorType::Rgb,
4 => ColorType::Rgba,
_ => return Err(PngError::InvalidChannelCount),
};
let mut enc = Encoder::new(&mut writer, w as u32, h as u32);
enc.set_color(colour);
enc.set_depth(BitDepth::Eight);
let mut whdr = enc.write_header()?;
let mut bytes = Vec::with_capacity(w * h * N);
bytes.extend(self.iter().flat_map(|px| px.to_bytes()));
whdr.write_image_data(&bytes)?;
Ok(())
}
}
fn match_colour_types(actual: ColorType, expected: ColorType) -> bool {
if actual == expected {
return true;
}
match (actual, expected) {
(ColorType::Rgb, ColorType::Rgba) => true,
(ColorType::Grayscale, ColorType::GrayscaleAlpha) => true,
(ColorType::Rgba, ColorType::Rgb) => true,
(ColorType::GrayscaleAlpha, ColorType::Grayscale) => true,
_ => false,
}
}