Do not accept images that have wrong sizes

This commit is contained in:
Jaby 2023-10-12 11:06:11 +02:00 committed by Jaby
parent f3d0fe17e0
commit c2955698e5
1 changed files with 171 additions and 147 deletions

View File

@ -1,148 +1,172 @@
use clap::{Args, ValueEnum}; use clap::{Args, ValueEnum};
use image::{DynamicImage, io::Reader as ImageReader}; use image::{DynamicImage, io::Reader as ImageReader};
use color_clut::IndexedImage; use color_clut::IndexedImage;
use color_full16::{RgbaImage, RgbImage}; use color_full16::{RgbaImage, RgbImage};
use std::io::{Cursor, Write}; use std::io::{Cursor, Write};
use tool_helper::{Error, Input}; use tool_helper::{Error, Input};
use types::{Header, Color as PSXColor, PSXImageConverter}; use types::{Header, Color as PSXColor, PSXImageConverter};
mod types; mod types;
mod color_clut; mod color_clut;
mod color_full16; mod color_full16;
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
pub enum ColorType{ pub enum ColorType{
Clut4, Clut4,
Clut8, Clut8,
Full16, Full16,
} }
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
pub enum ClutAlignment { pub enum ClutAlignment {
None, None,
Linear, Linear,
Block Block
} }
#[derive(Args)] #[derive(Args)]
pub struct Arguments { pub struct Arguments {
#[clap(value_enum, value_parser)] #[clap(value_enum, value_parser)]
color_depth: ColorType, color_depth: ColorType,
#[clap(value_enum, value_parser, default_value_t=ClutAlignment::None)] #[clap(value_enum, value_parser, default_value_t=ClutAlignment::None)]
clut_align: ClutAlignment, clut_align: ClutAlignment,
#[clap(long="semi-trans", default_value_t=false)] #[clap(long="semi-trans", default_value_t=false)]
semi_transparent: bool, semi_transparent: bool,
#[clap(long="color-trans", default_value_t=false)] #[clap(long="color-trans", default_value_t=false)]
transparent_palette: bool transparent_palette: bool
} }
fn modify_palette(mut image: IndexedImage, clut_align: ClutAlignment, semi_transparent: bool, transparent_palette: bool) -> IndexedImage { fn modify_palette(mut image: IndexedImage, clut_align: ClutAlignment, semi_transparent: bool, transparent_palette: bool) -> IndexedImage {
if semi_transparent { if semi_transparent {
for color in image.palette.iter_mut() { for color in image.palette.iter_mut() {
*color = PSXColor::semi_transparent(color.get_red(), color.get_green(), color.get_blue()); *color = PSXColor::semi_transparent(color.get_red(), color.get_green(), color.get_blue());
} }
} }
if transparent_palette { if transparent_palette {
if clut_align == ClutAlignment::Block { if clut_align == ClutAlignment::Block {
for color in image.palette.iter_mut().step_by(16) { for color in image.palette.iter_mut().step_by(16) {
*color = PSXColor::transparent(); *color = PSXColor::transparent();
} }
} }
else { else {
if let Some(first_color) = image.palette.get_mut(0) { if let Some(first_color) = image.palette.get_mut(0) {
*first_color = PSXColor::transparent(); *first_color = PSXColor::transparent();
} }
} }
} }
image image
} }
fn encode<T: PSXImageConverter>(image: T, color_depth: ColorType, clut_align: ClutAlignment, output: &mut dyn Write) -> Result<(), Error> { fn encode<T: PSXImageConverter>(image: T, color_depth: ColorType, clut_align: ClutAlignment, output: &mut dyn Write) -> Result<(), Error> {
let width = { let (width, height) = {
let width = image.width(); fn return_error(clut_type: u32, div: u32, width: u16, height: u16) -> Result<(u16, u16), Error> {
match color_depth { return Err(Error::from_callback(|| {format!("CLUT {} images require a width divideable by {} (found width: {}/{}={}, height: {}/{}={})", clut_type, div,
ColorType::Clut4 => width/4, width, div, (width as f32/div as f32),
ColorType::Clut8 => width/2, height, div, (height as f32/div as f32))}));
ColorType::Full16 => width }
}
}; let width = image.width();
let height = image.height(); let height = image.height();
let palette = image.get_palette(); match color_depth {
let (pal_width, pal_height) = { ColorType::Clut4 => {
if let Some(palette) = &palette { if width & 3 == 0 {
match clut_align { Ok((width/4, height))
ClutAlignment::None | }
ClutAlignment::Linear => (palette.len() as u16, 1u16),
ClutAlignment::Block => (16u16, (palette.len()/16) as u16), else {
} return_error(4, 4, width, height)
} }
},
else { ColorType::Clut8 => {
(0u16, 0u16) if width & 1 == 0 {
} Ok((width/2, height))
}; }
let header = Header::encode(width, height, pal_width, pal_height).ok_or(Error::from_callback(|| {format!("Image size (width: {}, height: {}) needs to be even", width, height)}))?; else {
return_error(8, 2, width, height)
tool_helper::raw::write_raw(output, &header)?; }
if let Some(palette) = palette { },
let mut color_count = pal_width*pal_height; ColorType::Full16 => {
for color in palette { Ok((width, height))
tool_helper::raw::write_raw(output, color)?; }
color_count -= 1; }
} }?;
let palette = image.get_palette();
while color_count > 0 { let (pal_width, pal_height) = {
tool_helper::raw::write_raw(output, &PSXColor::black())?; if let Some(palette) = &palette {
color_count -= 1; match clut_align {
} ClutAlignment::None |
} ClutAlignment::Linear => (palette.len() as u16, 1u16),
ClutAlignment::Block => (16u16, (palette.len()/16) as u16),
for color in image { }
tool_helper::raw::write_raw(output, &color)?; }
}
else {
Ok(()) (0u16, 0u16)
} }
};
fn convert_full16(input: Input, output: &mut dyn Write) -> Result<(), Error> {
match ImageReader::new(Cursor::new(tool_helper::input_to_vec(input)?)).with_guessed_format()?.decode() { let header = Header::encode(width, height, pal_width, pal_height).ok_or(Error::from_callback(|| {format!("Image size (width: {}, height: {}) needs to be even", width, height)}))?;
Ok(image) => {
match image { tool_helper::raw::write_raw(output, &header)?;
DynamicImage::ImageRgb8(image) => encode(RgbImage::new(image), ColorType::Full16, ClutAlignment::None, output), if let Some(palette) = palette {
DynamicImage::ImageRgba8(image) => encode(RgbaImage::new(image), ColorType::Full16, ClutAlignment::None, output), let mut color_count = pal_width*pal_height;
for color in palette {
_ => Err(Error::from_str("Only RGB and RGBA images are supported for 16bit encoding")) tool_helper::raw::write_raw(output, color)?;
} color_count -= 1;
}, }
Err(error) => Err(Error::from_error(error))
} while color_count > 0 {
} tool_helper::raw::write_raw(output, &PSXColor::black())?;
color_count -= 1;
fn convert_palette_based(input: Input, output: &mut dyn Write, color_type: ColorType, clut_align: ClutAlignment, semi_transparent: bool, transparent_palette: bool) -> Result<(), Error> { }
match png::Decoder::new(input).read_info() { }
Ok(reader) => {
let output_type = { for color in image {
match color_type { tool_helper::raw::write_raw(output, &color)?;
ColorType::Clut4 => color_clut::OutputType::FourBit, }
ColorType::Clut8 => color_clut::OutputType::EightBit,
_ => return Err(Error::from_str("ColorType not supported")) Ok(())
} }
};
fn convert_full16(input: Input, output: &mut dyn Write) -> Result<(), Error> {
encode(modify_palette(IndexedImage::new(reader, output_type)?, clut_align, semi_transparent, transparent_palette), color_type, clut_align, output) match ImageReader::new(Cursor::new(tool_helper::input_to_vec(input)?)).with_guessed_format()?.decode() {
}, Ok(image) => {
Err(error) => Err(Error::from_error(error)) match image {
} DynamicImage::ImageRgb8(image) => encode(RgbImage::new(image), ColorType::Full16, ClutAlignment::None, output),
} DynamicImage::ImageRgba8(image) => encode(RgbaImage::new(image), ColorType::Full16, ClutAlignment::None, output),
pub fn convert(args: Arguments, input: Input, output: &mut dyn Write) -> Result<(), Error> { _ => Err(Error::from_str("Only RGB and RGBA images are supported for 16bit encoding"))
match args.color_depth { }
ColorType::Full16 => convert_full16(input, output), },
_ => convert_palette_based(input, output, args.color_depth, args.clut_align, args.semi_transparent, args.transparent_palette), Err(error) => Err(Error::from_error(error))
} }
}
fn convert_palette_based(input: Input, output: &mut dyn Write, color_type: ColorType, clut_align: ClutAlignment, semi_transparent: bool, transparent_palette: bool) -> Result<(), Error> {
match png::Decoder::new(input).read_info() {
Ok(reader) => {
let output_type = {
match color_type {
ColorType::Clut4 => color_clut::OutputType::FourBit,
ColorType::Clut8 => color_clut::OutputType::EightBit,
_ => return Err(Error::from_str("ColorType not supported"))
}
};
encode(modify_palette(IndexedImage::new(reader, output_type)?, clut_align, semi_transparent, transparent_palette), color_type, clut_align, output)
},
Err(error) => Err(Error::from_error(error))
}
}
pub fn convert(args: Arguments, input: Input, output: &mut dyn Write) -> Result<(), Error> {
match args.color_depth {
ColorType::Full16 => convert_full16(input, output),
_ => convert_palette_based(input, output, args.color_depth, args.clut_align, args.semi_transparent, args.transparent_palette),
}
} }