Add pimg extract support

Cargo.toml

@@ -36,7 +36,7 @@ circus = []
 escude = ["int-enum"]
 escude-arc = ["escude", "rand", "utils-bit-stream"]
 kirikiri = ["emote-psb", "fancy-regex", "flate2", "utils-escape"]
-kirikiri-img = ["kirikiri", "image", "libtlg-rs"]
+kirikiri-img = ["kirikiri", "emote-psb", "image", "libtlg-rs"]
 yaneurao = []
 yaneurao-itufuru = ["yaneurao"]
 # basic feature

src/ext/psb.rs

@@ -96,6 +96,24 @@ impl PsbValueFixed {
         self.set_str(&value);
     }
 
+    pub fn as_u8(&self) -> Option<u8> {
+        self.as_i64().map(|n| n.try_into().ok()).flatten()
+    }
+
+    pub fn as_u32(&self) -> Option<u32> {
+        self.as_i64().map(|n| n as u32)
+    }
+
+    pub fn as_i64(&self) -> Option<i64> {
+        match self {
+            PsbValueFixed::Number(n) => match n {
+                PsbNumber::Integer(n) => Some(*n),
+                _ => None,
+            },
+            _ => None,
+        }
+    }
+
     pub fn as_str(&self) -> Option<&str> {
         match self {
             PsbValueFixed::String(s) => Some(s.string()),
@@ -138,6 +156,13 @@ impl PsbValueFixed {
             _ => ListIterMut::empty(),
         }
     }
+
+    pub fn resource_id(&self) -> Option<u64> {
+        match self {
+            PsbValueFixed::Resource(r) => Some(r.resource_ref),
+            _ => None,
+        }
+    }
 }
 
 impl Index<usize> for PsbValueFixed {

src/scripts/cat_system/image/hg3.rs

@@ -390,39 +390,3 @@ impl<'a> Hg3Reader<'a> {
         Ok(img)
     }
 }
-
-fn draw_on_canvas(
-    img: ImageData,
-    canvas_width: u32,
-    canvas_height: u32,
-    offset_x: u32,
-    offset_y: u32,
-) -> Result<ImageData> {
-    let bytes_per_pixel = img.color_type.bpp(img.depth) as u32 / 8;
-    let mut canvas_data = vec![0u8; (canvas_width * canvas_height * bytes_per_pixel) as usize];
-    let canvas_stride = canvas_width * bytes_per_pixel;
-    let img_stride = img.width * bytes_per_pixel;
-
-    for y in 0..img.height {
-        let canvas_y = y + offset_y;
-        if canvas_y >= canvas_height {
-            continue;
-        }
-        let canvas_start = (canvas_y * canvas_stride + offset_x * bytes_per_pixel) as usize;
-        let img_start = (y * img_stride) as usize;
-        let copy_len = img_stride as usize;
-        if canvas_start + copy_len > canvas_data.len() {
-            continue;
-        }
-        canvas_data[canvas_start..canvas_start + copy_len]
-            .copy_from_slice(&img.data[img_start..img_start + copy_len]);
-    }
-
-    Ok(ImageData {
-        width: canvas_width,
-        height: canvas_height,
-        color_type: img.color_type,
-        depth: img.depth,
-        data: canvas_data,
-    })
-}

src/scripts/kirikiri/image/mod.rs

@@ -1 +1,2 @@
+pub mod pimg;
 pub mod tlg;

src/scripts/kirikiri/image/pimg.rs

@@ -0,0 +1,334 @@
+use crate::ext::io::*;
+use crate::ext::psb::*;
+use crate::scripts::base::*;
+use crate::try_option;
+use crate::types::*;
+use crate::utils::img::*;
+use anyhow::Result;
+use emote_psb::PsbReader;
+use libtlg_rs::*;
+use std::collections::HashMap;
+use std::io::{Read, Seek};
+use std::path::Path;
+
+#[derive(Debug)]
+pub struct PImgBuilder {}
+
+impl PImgBuilder {
+    pub const fn new() -> Self {
+        Self {}
+    }
+}
+
+impl ScriptBuilder for PImgBuilder {
+    fn default_encoding(&self) -> Encoding {
+        Encoding::Utf8
+    }
+
+    fn build_script(
+        &self,
+        buf: Vec<u8>,
+        filename: &str,
+        _encoding: Encoding,
+        _archive_encoding: Encoding,
+        config: &ExtraConfig,
+    ) -> Result<Box<dyn Script>> {
+        Ok(Box::new(PImg::new(MemReader::new(buf), filename, config)?))
+    }
+
+    fn build_script_from_file(
+        &self,
+        filename: &str,
+        _encoding: Encoding,
+        _archive_encoding: Encoding,
+        config: &ExtraConfig,
+    ) -> Result<Box<dyn Script>> {
+        if filename == "-" {
+            let data = crate::utils::files::read_file(filename)?;
+            Ok(Box::new(PImg::new(MemReader::new(data), filename, config)?))
+        } else {
+            let f = std::fs::File::open(filename)?;
+            let reader = std::io::BufReader::new(f);
+            Ok(Box::new(PImg::new(reader, filename, config)?))
+        }
+    }
+
+    fn build_script_from_reader(
+        &self,
+        reader: Box<dyn ReadSeek>,
+        filename: &str,
+        _encoding: Encoding,
+        _archive_encoding: Encoding,
+        config: &ExtraConfig,
+    ) -> Result<Box<dyn Script>> {
+        Ok(Box::new(PImg::new(reader, filename, config)?))
+    }
+
+    fn extensions(&self) -> &'static [&'static str] {
+        &["pimg"]
+    }
+
+    fn script_type(&self) -> &'static ScriptType {
+        &ScriptType::KirikiriPimg
+    }
+
+    fn is_this_format(&self, filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
+        if Path::new(filename)
+            .extension()
+            .map(|ext| ext.to_ascii_lowercase() == "pimg")
+            .unwrap_or(false)
+            && buf_len >= 4
+            && buf.starts_with(b"PSB\0")
+        {
+            return Some(255);
+        }
+        None
+    }
+
+    fn is_image(&self) -> bool {
+        true
+    }
+}
+
+#[derive(Debug)]
+pub struct PImg {
+    psb: VirtualPsbFixed,
+}
+
+impl PImg {
+    pub fn new<R: Read + Seek>(reader: R, filename: &str, _config: &ExtraConfig) -> Result<Self> {
+        let mut psb = PsbReader::open_psb(reader)
+            .map_err(|e| anyhow::anyhow!("Failed to open PSB from {}: {:?}", filename, e))?;
+        let psb = psb
+            .load()
+            .map_err(|e| anyhow::anyhow!("Failed to load PSB from {}: {:?}", filename, e))?
+            .to_psb_fixed();
+        Ok(Self { psb })
+    }
+
+    fn load_img(&self, layer_id: i64) -> Result<Tlg> {
+        let layer_id = layer_id as usize;
+        let psb = self.psb.root();
+        let reference = &psb[format!("{layer_id}.tlg")];
+        let resource_id = reference
+            .resource_id()
+            .ok_or_else(|| anyhow::anyhow!("Layer {layer_id} does not have a resource ID"))?
+            as usize;
+        if resource_id >= self.psb.resources().len() {
+            return Err(anyhow::anyhow!(
+                "Resource ID {resource_id} for layer {layer_id} is out of bounds"
+            ));
+        }
+        let resource = &self.psb.resources()[resource_id];
+        Ok(load_tlg(MemReaderRef::new(&resource))?)
+    }
+}
+
+impl Script for PImg {
+    fn default_output_script_type(&self) -> OutputScriptType {
+        OutputScriptType::Json
+    }
+
+    fn default_format_type(&self) -> FormatOptions {
+        FormatOptions::None
+    }
+
+    fn is_image(&self) -> bool {
+        true
+    }
+
+    fn is_multi_image(&self) -> bool {
+        true
+    }
+
+    fn export_multi_image<'a>(
+        &'a self,
+    ) -> Result<Box<dyn Iterator<Item = Result<ImageDataWithName>> + 'a>> {
+        let psb = self.psb.root();
+        let width = psb["width"]
+            .as_u32()
+            .ok_or(anyhow::anyhow!("missing width"))?;
+        let height = psb["height"]
+            .as_u32()
+            .ok_or(anyhow::anyhow!("missing height"))?;
+        if !psb["layers"].is_list() {
+            return Err(anyhow::anyhow!("layers is not a list"));
+        }
+        if psb["layers"].len() == 0 {
+            return Ok(Box::new(std::iter::empty()));
+        }
+        let mut bases = HashMap::new();
+        for i in psb["layers"].members() {
+            if !i["diff_id"].is_none() {
+                continue; // Skip layers with diff_id
+            }
+            let layer_id = i["layer_id"]
+                .as_i64()
+                .ok_or(anyhow::anyhow!("missing layer_id"))?;
+            let top = i["top"].as_u32().ok_or(anyhow::anyhow!("missing top"))?;
+            let left = i["left"].as_u32().ok_or(anyhow::anyhow!("missing left"))?;
+            let opacity = i["opacity"]
+                .as_u8()
+                .ok_or_else(|| anyhow::anyhow!("Layer does not have a valid opacity"))?;
+            bases.insert(layer_id, (self.load_img(layer_id)?, top, left, opacity));
+        }
+        Ok(Box::new(PImgIter {
+            pimg: self,
+            width,
+            height,
+            layers: psb["layers"].members(),
+            bases,
+        }))
+    }
+}
+
+struct PImgIter<'a> {
+    pimg: &'a PImg,
+    width: u32,
+    height: u32,
+    layers: ListIter<'a>,
+    bases: HashMap<i64, (Tlg, u32, u32, u8)>,
+}
+
+impl<'a> Iterator for PImgIter<'a> {
+    type Item = Result<ImageDataWithName>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.layers.next() {
+            Some(layer) => {
+                let layer_id =
+                    try_option!(layer["layer_id"].as_i64().ok_or_else(|| {
+                        anyhow::anyhow!("Layer does not have a valid layer_id")
+                    }));
+                let layer_name = try_option!(
+                    layer["name"]
+                        .as_str()
+                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid name") })
+                );
+                let width = try_option!(
+                    layer["width"]
+                        .as_u32()
+                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid width") })
+                );
+                let height = try_option!(
+                    layer["height"]
+                        .as_u32()
+                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid height") })
+                );
+                let top = try_option!(
+                    layer["top"]
+                        .as_u32()
+                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid top") })
+                );
+                let left = try_option!(
+                    layer["left"]
+                        .as_u32()
+                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid left") })
+                );
+                let opacity = try_option!(
+                    layer["opacity"]
+                        .as_u8()
+                        .ok_or_else(|| { anyhow::anyhow!("Layer does not have a valid opacity") })
+                );
+                if layer["diff_id"].is_none() {
+                    let base = &try_option!(self.bases.get(&layer_id).ok_or(anyhow::anyhow!(
+                        "Base image for layer_id {} not found",
+                        layer_id
+                    )))
+                    .0;
+                    let mut data = ImageData {
+                        width: self.width,
+                        height: self.height,
+                        color_type: match base.color {
+                            TlgColorType::Bgr24 => ImageColorType::Bgr,
+                            TlgColorType::Bgra32 => ImageColorType::Bgra,
+                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
+                        },
+                        depth: 8,
+                        data: base.data.clone(),
+                    };
+                    if opacity != 255 {
+                        try_option!(apply_opacity(&mut data, opacity));
+                    }
+                    if self.width != width || self.height != height || top != 0 || left != 0 {
+                        data =
+                            try_option!(draw_on_canvas(data, self.width, self.height, left, top));
+                    }
+                    return Some(Ok(ImageDataWithName {
+                        name: layer_name.to_string(),
+                        data,
+                    }));
+                } else {
+                    let diff_id =
+                        try_option!(layer["diff_id"].as_i64().ok_or_else(|| {
+                            anyhow::anyhow!("Layer does not have a valid diff_id")
+                        }));
+                    let (base, base_top, base_left, base_opacity) = try_option!(
+                        self.bases
+                            .get(&diff_id)
+                            .ok_or(anyhow::anyhow!("Base image layer {} not found", diff_id))
+                    );
+                    let diff = try_option!(self.pimg.load_img(layer_id));
+                    if base.color != diff.color {
+                        return Some(Err(anyhow::anyhow!(
+                            "Color type mismatch for layer_id {}: base color {:?}, diff color {:?}",
+                            layer_id,
+                            base.color,
+                            diff.color
+                        )));
+                    }
+                    let mut base_img = ImageData {
+                        width: base.width,
+                        height: base.height,
+                        color_type: match base.color {
+                            TlgColorType::Bgr24 => ImageColorType::Bgr,
+                            TlgColorType::Bgra32 => ImageColorType::Bgra,
+                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
+                        },
+                        depth: 8,
+                        data: base.data.clone(),
+                    };
+                    if base.width != self.width
+                        || base.height != self.height
+                        || *base_top != 0
+                        || *base_left != 0
+                    {
+                        base_img = try_option!(draw_on_canvas(
+                            base_img,
+                            self.width,
+                            self.height,
+                            *base_left,
+                            *base_top
+                        ));
+                    }
+                    if *base_opacity != 255 {
+                        try_option!(apply_opacity(&mut base_img, *base_opacity));
+                    }
+                    let diff = ImageData {
+                        width: diff.width,
+                        height: diff.height,
+                        color_type: match diff.color {
+                            TlgColorType::Bgr24 => ImageColorType::Bgr,
+                            TlgColorType::Bgra32 => ImageColorType::Bgra,
+                            TlgColorType::Grayscale8 => ImageColorType::Grayscale,
+                        },
+                        depth: 8,
+                        data: diff.data.clone(),
+                    };
+                    try_option!(draw_on_img_with_opacity(
+                        &mut base_img,
+                        &diff,
+                        left,
+                        top,
+                        opacity
+                    ));
+                    Some(Ok(ImageDataWithName {
+                        name: layer_name.to_string(),
+                        data: base_img,
+                    }))
+                }
+            }
+            None => None,
+        }
+    }
+}

src/scripts/mod.rs

@@ -56,6 +56,8 @@ lazy_static::lazy_static! {
         Box::new(kirikiri::ks::KsBuilder::new()),
         #[cfg(feature = "kirikiri-img")]
         Box::new(kirikiri::image::tlg::TlgImageBuilder::new()),
+        #[cfg(feature = "kirikiri-img")]
+        Box::new(kirikiri::image::pimg::PImgBuilder::new()),
     ];
     pub static ref ALL_EXTS: Vec<String> =
         BUILDER.iter().flat_map(|b| b.extensions()).map(|s| s.to_string()).collect();

src/types.rs

@@ -290,6 +290,10 @@ pub enum ScriptType {
     #[value(alias("kr-tlg"))]
     /// Kirikiri TLG image
     KirikiriTlg,
+    #[cfg(feature = "kirikiri-img")]
+    #[value(alias("kr-pimg"))]
+    /// Kirikiri PIMG image
+    KirikiriPimg,
     #[cfg(feature = "yaneurao-itufuru")]
     #[value(alias("itufuru"))]
     /// Yaneurao Itufuru script

src/utils/img.rs

@@ -161,6 +161,42 @@ pub fn decode_img(typ: ImageOutputType, filename: &str) -> Result<ImageData> {
     }
 }
 
+pub fn draw_on_canvas(
+    img: ImageData,
+    canvas_width: u32,
+    canvas_height: u32,
+    offset_x: u32,
+    offset_y: u32,
+) -> Result<ImageData> {
+    let bytes_per_pixel = img.color_type.bpp(img.depth) as u32 / 8;
+    let mut canvas_data = vec![0u8; (canvas_width * canvas_height * bytes_per_pixel) as usize];
+    let canvas_stride = canvas_width * bytes_per_pixel;
+    let img_stride = img.width * bytes_per_pixel;
+
+    for y in 0..img.height {
+        let canvas_y = y + offset_y;
+        if canvas_y >= canvas_height {
+            continue;
+        }
+        let canvas_start = (canvas_y * canvas_stride + offset_x * bytes_per_pixel) as usize;
+        let img_start = (y * img_stride) as usize;
+        let copy_len = img_stride as usize;
+        if canvas_start + copy_len > canvas_data.len() {
+            continue;
+        }
+        canvas_data[canvas_start..canvas_start + copy_len]
+            .copy_from_slice(&img.data[img_start..img_start + copy_len]);
+    }
+
+    Ok(ImageData {
+        width: canvas_width,
+        height: canvas_height,
+        color_type: img.color_type,
+        depth: img.depth,
+        data: canvas_data,
+    })
+}
+
 pub fn flip_image(data: &mut ImageData) -> Result<()> {
     if data.height <= 1 {
         return Ok(());
@@ -183,3 +219,93 @@ pub fn flip_image(data: &mut ImageData) -> Result<()> {
 
     Ok(())
 }
+
+pub fn apply_opacity(img: &mut ImageData, opacity: u8) -> Result<()> {
+    if img.color_type != ImageColorType::Rgba && img.color_type != ImageColorType::Bgra {
+        return Err(anyhow::anyhow!("Image is not RGBA or BGRA"));
+    }
+    if img.depth != 8 {
+        return Err(anyhow::anyhow!(
+            "Opacity application only supports 8-bit depth"
+        ));
+    }
+    for i in (0..img.data.len()).step_by(4) {
+        img.data[i + 3] = (img.data[i + 3] as u16 * opacity as u16 / 255) as u8;
+    }
+    Ok(())
+}
+
+pub fn draw_on_img_with_opacity(
+    base: &mut ImageData,
+    diff: &ImageData,
+    left: u32,
+    top: u32,
+    opacity: u8,
+) -> Result<()> {
+    if base.color_type != diff.color_type {
+        return Err(anyhow::anyhow!("Image color types do not match"));
+    }
+    if base.color_type != ImageColorType::Rgba && base.color_type != ImageColorType::Bgra {
+        return Err(anyhow::anyhow!("Images are not RGBA or BGRA"));
+    }
+    if base.depth != 8 || diff.depth != 8 {
+        return Err(anyhow::anyhow!(
+            "Image drawing with opacity only supports 8-bit depth"
+        ));
+    }
+
+    let bpp = 4;
+    let base_stride = base.width as usize * bpp;
+    let diff_stride = diff.width as usize * bpp;
+
+    for y in 0..diff.height {
+        let base_y = top + y;
+        if base_y >= base.height {
+            continue;
+        }
+
+        for x in 0..diff.width {
+            let base_x = left + x;
+            if base_x >= base.width {
+                continue;
+            }
+
+            let diff_idx = (y as usize * diff_stride) + (x as usize * bpp);
+            let base_idx = (base_y as usize * base_stride) + (base_x as usize * bpp);
+
+            let diff_pixel = &diff.data[diff_idx..diff_idx + bpp];
+            let base_pixel_orig = base.data[base_idx..base_idx + bpp].to_vec();
+
+            let src_alpha_u16 = (diff_pixel[3] as u16 * opacity as u16) / 255;
+
+            if src_alpha_u16 == 0 {
+                continue;
+            }
+
+            let dst_alpha_u16 = base_pixel_orig[3] as u16;
+
+            // out_alpha = src_alpha + dst_alpha * (1 - src_alpha)
+            let out_alpha_u16 = src_alpha_u16 + (dst_alpha_u16 * (255 - src_alpha_u16)) / 255;
+
+            if out_alpha_u16 == 0 {
+                for i in 0..4 {
+                    base.data[base_idx + i] = 0;
+                }
+                continue;
+            }
+
+            // out_color = (src_color * src_alpha + dst_color * dst_alpha * (1 - src_alpha)) / out_alpha
+            for i in 0..3 {
+                let src_comp = diff_pixel[i] as u16;
+                let dst_comp = base_pixel_orig[i] as u16;
+
+                let numerator = src_comp * src_alpha_u16
+                    + (dst_comp * dst_alpha_u16 * (255 - src_alpha_u16)) / 255;
+                base.data[base_idx + i] = (numerator / out_alpha_u16) as u8;
+            }
+            base.data[base_idx + 3] = out_alpha_u16 as u8;
+        }
+    }
+
+    Ok(())
+}

src/utils/macros.rs

@@ -0,0 +1,13 @@
+#[macro_export]
+macro_rules! try_option {
+    ($expr:expr $(,)?) => {
+        match $expr {
+            std::result::Result::Ok(val) => val,
+            std::result::Result::Err(err) => {
+                return std::option::Option::Some(std::result::Result::Err(
+                    std::convert::From::from(err),
+                ));
+            }
+        }
+    };
+}

src/utils/mod.rs

@@ -11,5 +11,6 @@ pub mod escape;
 pub mod files;
 #[cfg(feature = "image")]
 pub mod img;
+pub mod macros;
 pub mod name_replacement;
 pub mod struct_pack;