msg_tool\utils\psd/
mod.rs

1//! A simple PSD reader/writer
2mod compression;
3mod types;
4
5use crate::ext::io::*;
6use crate::types::*;
7use crate::utils::bit_stream::*;
8use crate::utils::encoding::*;
9use crate::utils::img::*;
10use crate::utils::struct_pack::*;
11use anyhow::Result;
12use compression::*;
13use std::collections::BTreeMap;
14use std::io::{Read, Seek, Write};
15use std::ops::Deref;
16use types::*;
17
18pub use types::{
19    AdditionalLayerInfo, IMAGE_RESOURCE_SIGNATURE, LAYER_ID_KEY, LAYER_NAME_SOURCE_SETTING_KEY,
20    LayerID, LayerNameSourceSetting,
21};
22
23#[derive(Debug, Clone, msg_tool_macro::Default)]
24pub struct PsdLayerOption {
25    #[default(true)]
26    /// Whether the layer is visible.
27    pub visible: bool,
28    #[default(255)]
29    /// The opacity of the layer (0-255).
30    pub opacity: u8,
31    /// Additional layer information.
32    pub additional_info: Vec<AdditionalLayerInfo>,
33}
34
35impl PsdLayerOption {
36    fn to_flags(&self) -> u8 {
37        let mut flags = 0u8;
38        if !self.visible {
39            flags |= 0b0000_0010;
40        }
41        flags
42    }
43}
44
45/// A simple PSD writer.
46pub struct PsdWriter {
47    psd: PsdFile,
48    color_type: ImageColorType,
49    compress: bool,
50    zlib_compression_level: u32,
51    encoding: Encoding,
52}
53
54fn encode_unicode_layer(name: &str) -> Result<AdditionalLayerInfo> {
55    let layer = UnicodeLayer {
56        name: UnicodeString(name.to_string()),
57    };
58    let mut data = MemWriter::new();
59    layer.pack(&mut data, true, Encoding::Utf16BE, &None)?;
60    Ok(AdditionalLayerInfo {
61        signature: *IMAGE_RESOURCE_SIGNATURE,
62        key: *UNICODE_LAYER_KEY,
63        data: data.into_inner(),
64    })
65}
66
67impl PsdWriter {
68    /// Creates a new PSD writer with the specified dimensions, color type, and bit depth.
69    pub fn new(
70        width: u32,
71        height: u32,
72        color_type: ImageColorType,
73        depth: u8,
74        encoding: Encoding,
75    ) -> Result<Self> {
76        let color_type = match color_type {
77            ImageColorType::Bgr => ImageColorType::Rgb,
78            ImageColorType::Bgra => ImageColorType::Rgba,
79            _ => color_type,
80        };
81        let depth = match depth {
82            1 | 8 | 16 | 32 => depth,
83            _ => anyhow::bail!("Unsupported bit depth: {}", depth),
84        };
85        let psd = PsdFile {
86            header: PsdHeader {
87                signature: *PSD_SIGNATURE,
88                version: 1,
89                reserved: [0; 6],
90                channels: color_type.bpp(1),
91                height,
92                width,
93                depth: depth as u16,
94                color_mode: if color_type == ImageColorType::Grayscale {
95                    1
96                } else {
97                    3
98                },
99            },
100            color_mode_data: ColorModeData { data: vec![] },
101            image_resource: ImageResourceSection { resources: vec![] },
102            layer_and_mask_info: LayerAndMaskInfo {
103                layer_info: LayerInfo {
104                    layer_count: 0,
105                    layer_records: vec![],
106                    channel_image_data: vec![],
107                },
108                global_layer_mask_info: Some(GlobalLayerMaskInfo {
109                    overlays_color_space: 0,
110                    overlays_color_components: [0; 4],
111                    opacity: 0,
112                    kind: 128,
113                    filler: vec![0],
114                }),
115                tagged_blocks: vec![],
116            },
117            image_data: ImageDataSection {
118                compression: 0,
119                image_data: vec![],
120            },
121        };
122        Ok(Self {
123            psd,
124            color_type,
125            compress: true,
126            zlib_compression_level: 6,
127            encoding,
128        })
129    }
130
131    /// Sets whether to compress image data in the PSD file.
132    pub fn compress(mut self, compress: bool) -> Self {
133        self.compress = compress;
134        self
135    }
136
137    /// Sets the zlib compression level for the PSD file.
138    pub fn zlib_compression_level(mut self, level: u32) -> Self {
139        self.zlib_compression_level = level;
140        self
141    }
142
143    /// Add a layer to the PSD file.
144    ///
145    /// * `name` - The name of the layer.
146    /// * `x` - The x position of the layer.
147    /// * `y` - The y position of the layer.
148    /// * `data` - The image data of the layer.
149    /// * `option` - The options for the layer.
150    pub fn add_layer(
151        &mut self,
152        name: &str,
153        x: u32,
154        y: u32,
155        mut data: ImageData,
156        option: Option<PsdLayerOption>,
157    ) -> Result<()> {
158        if data.color_type == ImageColorType::Bgr {
159            convert_bgr_to_rgb(&mut data)?;
160        }
161        if data.color_type == ImageColorType::Bgra {
162            convert_bgra_to_rgba(&mut data)?;
163        }
164        let length = data.width as u32 * data.height as u32;
165        let mut channel_ids = Vec::new();
166        if data.color_type == ImageColorType::Grayscale {
167            channel_ids.push(0);
168        } else {
169            channel_ids.push(0); // R
170            channel_ids.push(1); // G
171            channel_ids.push(2); // B
172            if data.color_type == ImageColorType::Rgba {
173                channel_ids.push(-1); // Alpha
174            }
175        }
176        let flags = if let Some(opt) = &option {
177            opt.to_flags()
178        } else {
179            0
180        };
181        let opacity = if let Some(opt) = &option {
182            opt.opacity
183        } else {
184            255
185        };
186        let mut layer_base = LayerRecordBase {
187            top: y as i32,
188            left: x as i32,
189            bottom: (y + data.height) as i32,
190            right: (x + data.width) as i32,
191            channels: data.color_type.bpp(1) as u16,
192            channel_infos: Vec::new(),
193            blend_mode_signature: *IMAGE_RESOURCE_SIGNATURE,
194            blend_mode_key: *b"norm",
195            opacity,
196            clipping: 0,
197            flags,
198            filler: 0,
199        };
200        let mut channel_ranges = Vec::new();
201        for _ in 0..layer_base.channels {
202            channel_ranges.push(ChannelRange {
203                source_range: 0xFFFF,
204                dest_range: 0xFFFF,
205            });
206        }
207        let layer_blending_ranges = LayerBlendingRanges {
208            gray_blend_dest: 0xFFFF,
209            gray_blend_source: 0xFFFF,
210            channel_ranges,
211        };
212        let mut image_data = Vec::new();
213        for i in 0..layer_base.channels {
214            let mut d = Vec::with_capacity(length as usize);
215            for y in 0..data.height {
216                for x in 0..data.width {
217                    let index =
218                        (y * data.width + x) as usize * layer_base.channels as usize + i as usize;
219                    d.push(data.data[index]);
220                }
221            }
222            if self.compress {
223                for y in 0..data.height {
224                    let ind = y as usize * data.width as usize;
225                    let mut pre = d[ind];
226                    for x in 1..data.width as usize {
227                        let cur = d[ind + x];
228                        d[ind + x] = cur.wrapping_sub(pre);
229                        pre = cur;
230                    }
231                }
232                let mut data = Vec::new();
233                let mut enc = flate2::write::ZlibEncoder::new(
234                    &mut data,
235                    flate2::Compression::new(self.zlib_compression_level),
236                );
237                enc.write_all(&d)?;
238                enc.finish()?;
239                d = data;
240            }
241            let cinfo = ChannelInfo {
242                channel_id: channel_ids[i as usize],
243                length: d.len() as u32 + 2, // +2 for compression method
244            };
245            layer_base.channel_infos.push(cinfo);
246            let compression = if self.compress { 3 } else { 0 };
247            image_data.push(ChannelImageData {
248                compression,
249                image_data: d,
250            });
251        }
252        let encoded = encode_string(self.encoding, &name, false)?;
253        let mut infos = vec![encode_unicode_layer(name)?];
254        if let Some(opt) = option {
255            infos.extend(opt.additional_info);
256        }
257        let layer = LayerRecord {
258            base: layer_base,
259            layer_mask: None,
260            layer_blending_ranges,
261            layer_name: PascalString4(encoded),
262            infos,
263        };
264        self.psd
265            .layer_and_mask_info
266            .layer_info
267            .layer_records
268            .push(layer);
269
270        // Update layer count
271        self.psd.layer_and_mask_info.layer_info.layer_count += 1;
272
273        self.psd
274            .layer_and_mask_info
275            .layer_info
276            .channel_image_data
277            .extend(image_data);
278        Ok(())
279    }
280
281    /// Adds the start of a layer group to the PSD file.
282    pub fn add_layer_group(
283        &mut self,
284        name: &str,
285        is_closed: bool,
286        option: Option<PsdLayerOption>,
287    ) -> Result<()> {
288        let type_info = SectionDividerSetting {
289            typ: if is_closed { 2 } else { 1 },
290        };
291        let mut data = MemWriter::new();
292        type_info.pack(&mut data, true, self.encoding, &None)?;
293        let encoded = encode_string(self.encoding, &name, false)?;
294        let flags = if let Some(opt) = &option {
295            opt.to_flags()
296        } else {
297            0
298        };
299        let opacity = if let Some(opt) = &option {
300            opt.opacity
301        } else {
302            255
303        };
304        let mut infos = vec![
305            AdditionalLayerInfo {
306                signature: *IMAGE_RESOURCE_SIGNATURE,
307                key: *SECTION_DIVIDER_SETTING_KEY,
308                data: data.into_inner(),
309            },
310            encode_unicode_layer(name)?,
311        ];
312        if let Some(opt) = option {
313            infos.extend(opt.additional_info);
314        }
315        let layer = LayerRecord {
316            base: LayerRecordBase {
317                top: 0,
318                left: 0,
319                bottom: 0,
320                right: 0,
321                channels: 0,
322                channel_infos: vec![],
323                blend_mode_signature: *IMAGE_RESOURCE_SIGNATURE,
324                blend_mode_key: *b"pass",
325                opacity,
326                clipping: 0,
327                flags,
328                filler: 0,
329            },
330            layer_mask: None,
331            layer_blending_ranges: LayerBlendingRanges {
332                gray_blend_dest: 0xFFFF,
333                gray_blend_source: 0xFFFF,
334                channel_ranges: vec![],
335            },
336            layer_name: PascalString4(encoded),
337            infos,
338        };
339        self.psd
340            .layer_and_mask_info
341            .layer_info
342            .layer_records
343            .push(layer);
344        self.psd.layer_and_mask_info.layer_info.layer_count += 1;
345        Ok(())
346    }
347
348    /// Adds the end of a layer group to the PSD file.
349    pub fn add_layer_group_end(&mut self) -> Result<()> {
350        let type_info = SectionDividerSetting { typ: 3 };
351        let mut data = MemWriter::new();
352        type_info.pack(&mut data, true, self.encoding, &None)?;
353        let layer = LayerRecord {
354            base: LayerRecordBase {
355                top: 0,
356                left: 0,
357                bottom: 0,
358                right: 0,
359                channels: 0,
360                channel_infos: vec![],
361                blend_mode_signature: *IMAGE_RESOURCE_SIGNATURE,
362                blend_mode_key: *b"norm",
363                opacity: 255,
364                clipping: 0,
365                flags: 0,
366                filler: 0,
367            },
368            layer_mask: None,
369            layer_blending_ranges: LayerBlendingRanges {
370                gray_blend_dest: 0xFFFF,
371                gray_blend_source: 0xFFFF,
372                channel_ranges: vec![],
373            },
374            layer_name: PascalString4(b"</Layer group>".to_vec()),
375            infos: vec![AdditionalLayerInfo {
376                signature: *IMAGE_RESOURCE_SIGNATURE,
377                key: *SECTION_DIVIDER_SETTING_KEY,
378                data: data.into_inner(),
379            }],
380        };
381        self.psd
382            .layer_and_mask_info
383            .layer_info
384            .layer_records
385            .push(layer);
386        self.psd.layer_and_mask_info.layer_info.layer_count += 1;
387        Ok(())
388    }
389
390    /// Saves the PSD file to the specified writer with the given encoding.
391    ///
392    /// * `data` - The final composite image data to be saved in the PSD file.
393    pub fn save<T: Write>(&mut self, data: ImageData, mut writer: T) -> Result<()> {
394        if data.color_type == ImageColorType::Bgr {
395            convert_bgr_to_rgb(&mut data.clone())?;
396        }
397        if data.color_type == ImageColorType::Bgra {
398            convert_bgra_to_rgba(&mut data.clone())?;
399        }
400        if self.color_type != data.color_type {
401            anyhow::bail!(
402                "Image color type does not match PSD color type: {:?} != {:?}",
403                self.color_type,
404                data.color_type
405            );
406        }
407        if data.width != self.psd.header.width || data.height != self.psd.header.height {
408            anyhow::bail!(
409                "Image dimensions do not match PSD dimensions: {}x{} != {}x{}",
410                data.width,
411                data.height,
412                self.psd.header.width,
413                self.psd.header.height
414            );
415        }
416
417        // Convert interleaved data to planar data (RRR...GGG...BBB...)
418        let channels = self.psd.header.channels as usize;
419        let width = self.psd.header.width as usize;
420        let height = self.psd.header.height as usize;
421        let total_pixels = width * height;
422        let expected_len = total_pixels * channels; // assuming 8-bit depth
423
424        if data.data.len() != expected_len {
425            anyhow::bail!("Data length mismatch for planar conversion");
426        }
427
428        let mut planar_data = Vec::with_capacity(expected_len);
429        for c in 0..channels {
430            for i in 0..total_pixels {
431                // Interleaved index: pixel_index * channels + channel_index
432                let val = data.data[i * channels + c];
433                planar_data.push(val);
434            }
435        }
436        if self.compress {
437            // RLE compression for planar data
438            let mut compressed = MemWriter::new();
439            // reserve 2 bytes per scanline for lengths
440            for _ in 0..(channels * height) {
441                compressed.write_u16_be(0)?; // placeholder for lengths
442            }
443            for c in 0..channels {
444                for y in 0..height {
445                    let start = (c * width * height) + (y * width);
446                    let line_end = start + width;
447                    let out_line = rle_compress(&planar_data[start..line_end]);
448
449                    // write scanline length at reserved spot and append data
450                    compressed
451                        .write_u16_be_at(((c * height + y) * 2) as u64, out_line.len() as u16)?;
452                    compressed.write_all(&out_line)?;
453                }
454            }
455            planar_data = compressed.into_inner();
456        }
457        let compression = if self.compress { 1 } else { 0 };
458        self.psd.image_data.image_data = planar_data;
459        self.psd.image_data.compression = compression;
460        self.psd.pack(&mut writer, true, self.encoding, &None)?;
461        Ok(())
462    }
463}
464
465#[derive(Debug)]
466pub struct PsdReader {
467    psd: PsdFile,
468    encoding: Encoding,
469    channel_start_indices: Vec<usize>,
470}
471
472#[derive(Debug)]
473pub struct NormalLayer<'a> {
474    layer: &'a LayerRecord,
475    layer_idx: usize,
476    psd: &'a PsdReader,
477}
478
479impl<'a> NormalLayer<'a> {
480    /// Returns the name of the layer.
481    pub fn layer_name(&self) -> Result<String> {
482        self.layer.layer_name(self.psd.encoding)
483    }
484
485    /// Returns the current layer's index in the PSD file.
486    pub fn layer_index(&self) -> usize {
487        self.layer_idx
488    }
489
490    /// Returns the top position of the layer.
491    pub fn top(&self) -> i32 {
492        self.layer.base.top
493    }
494
495    /// Returns the left position of the layer.
496    pub fn left(&self) -> i32 {
497        self.layer.base.left
498    }
499
500    /// Returns the bottom position of the layer.
501    pub fn bottom(&self) -> i32 {
502        self.layer.base.bottom
503    }
504
505    /// Returns the right position of the layer.
506    pub fn right(&self) -> i32 {
507        self.layer.base.right
508    }
509
510    /// Returns the width of the layer.
511    pub fn width(&self) -> u32 {
512        (self.layer.base.right - self.layer.base.left) as u32
513    }
514
515    /// Returns the height of the layer.
516    pub fn height(&self) -> u32 {
517        (self.layer.base.bottom - self.layer.base.top) as u32
518    }
519
520    /// Returns the number of channels in the layer.
521    pub fn channels(&self) -> u16 {
522        self.layer.base.channels
523    }
524
525    /// Reads and returns the raw channel id and data of the layer.
526    pub fn read_raw_data(&self) -> Result<Vec<(i16, Vec<u8>)>> {
527        let mut start_idx = self.psd.channel_start_indices[self.layer_idx];
528        let mut channels_data = Vec::new();
529        for cinfo in &self.layer.base.channel_infos {
530            let mut data = self
531                .psd
532                .psd
533                .layer_and_mask_info
534                .layer_info
535                .channel_image_data[start_idx]
536                .clone();
537            start_idx += 1;
538            decompress_channel_image_data(&mut data, self)?;
539            channels_data.push((cinfo.channel_id, data.image_data));
540        }
541        Ok(channels_data)
542    }
543
544    /// Reads and returns the image data of the layer.
545    pub fn image(&self) -> Result<ImageData> {
546        let color_mode = self.psd.color_mode();
547        if !matches!(color_mode, 1 | 3) {
548            anyhow::bail!(
549                "Unsupported PSD color mode for image extraction: {}",
550                self.psd.color_mode()
551            );
552        }
553        let channels = self.layer.base.channels;
554        let width = (self.layer.base.right - self.layer.base.left) as u32;
555        let height = (self.layer.base.bottom - self.layer.base.top) as u32;
556        if channels < 1 {
557            anyhow::bail!("PSD layer has no channels");
558        }
559        let channels_map = BTreeMap::from_iter(self.read_raw_data()?);
560        if color_mode == 1 {
561            let grayscale = channels_map
562                .get(&0)
563                .ok_or_else(|| anyhow::anyhow!("PSD grayscale layer missing channel 0"))?;
564            let alpha = channels_map.get(&-1);
565            let depth = self.psd.bit_depth() as u8;
566            if let Some(alpha) = alpha {
567                let mut g = MsbBitStream::new(MemReaderRef::new(&grayscale));
568                let mut a = MsbBitStream::new(MemReaderRef::new(&alpha));
569                let mut data = MemWriter::new();
570                let mut o = MsbBitWriter::new(&mut data);
571                for _ in 0..height {
572                    g.m_cached_bits = 0;
573                    a.m_cached_bits = 0;
574                    for _ in 0..width {
575                        let gray = g.get_bits(depth as u32)?;
576                        let alpha = a.get_bits(depth as u32)?;
577                        o.put_bits(gray, depth)?;
578                        o.put_bits(gray, depth)?;
579                        o.put_bits(gray, depth)?;
580                        o.put_bits(alpha, depth)?;
581                    }
582                    o.flush()?;
583                }
584                Ok(ImageData {
585                    width,
586                    height,
587                    color_type: ImageColorType::Rgba,
588                    depth,
589                    data: data.into_inner(),
590                })
591            } else {
592                Ok(ImageData {
593                    width,
594                    height,
595                    color_type: ImageColorType::Grayscale,
596                    depth,
597                    data: grayscale.clone(),
598                })
599            }
600        } else {
601            let red = channels_map
602                .get(&0)
603                .ok_or_else(|| anyhow::anyhow!("PSD RGB layer missing channel 0"))?;
604            let green = channels_map
605                .get(&1)
606                .ok_or_else(|| anyhow::anyhow!("PSD RGB layer missing channel 1"))?;
607            let blue = channels_map
608                .get(&2)
609                .ok_or_else(|| anyhow::anyhow!("PSD RGB layer missing channel 2"))?;
610            let mut a = channels_map
611                .get(&-1)
612                .map(|v| MsbBitStream::new(MemReaderRef::new(v)));
613            let depth = self.psd.bit_depth() as u8;
614            let mut r = MsbBitStream::new(MemReaderRef::new(&red));
615            let mut g = MsbBitStream::new(MemReaderRef::new(&green));
616            let mut b = MsbBitStream::new(MemReaderRef::new(&blue));
617            let mut data = MemWriter::new();
618            let mut o = MsbBitWriter::new(&mut data);
619            for _ in 0..height {
620                r.m_cached_bits = 0;
621                g.m_cached_bits = 0;
622                b.m_cached_bits = 0;
623                if let Some(alpha) = &mut a {
624                    alpha.m_cached_bits = 0;
625                }
626                for _ in 0..width {
627                    let red = r.get_bits(depth as u32)?;
628                    let green = g.get_bits(depth as u32)?;
629                    let blue = b.get_bits(depth as u32)?;
630                    o.put_bits(red, depth)?;
631                    o.put_bits(green, depth)?;
632                    o.put_bits(blue, depth)?;
633                    if let Some(alpha) = &mut a {
634                        let alpha = alpha.get_bits(depth as u32)?;
635                        o.put_bits(alpha, depth)?;
636                    }
637                }
638                o.flush()?;
639            }
640            Ok(ImageData {
641                width,
642                height,
643                color_type: if a.is_some() {
644                    ImageColorType::Rgba
645                } else {
646                    ImageColorType::Rgb
647                },
648                depth,
649                data: data.into_inner(),
650            })
651        }
652    }
653}
654
655#[derive(Debug)]
656pub struct GroupLayer<'a> {
657    layer: &'a LayerRecord,
658    layer_idx: usize,
659    psd: &'a PsdReader,
660    pub childrens: Vec<Layer<'a>>,
661    is_closed: bool,
662}
663
664impl<'a> GroupLayer<'a> {
665    fn create<'b>(layer_idx: &'b mut usize, psd: &'a PsdReader) -> Result<Self> {
666        let mut childrens = Vec::new();
667        // skip the end marker
668        *layer_idx += 1;
669        let layer_count = psd.psd.layer_count();
670        while *layer_idx < layer_count {
671            let layer = &psd.psd.layer_and_mask_info.layer_info.layer_records[*layer_idx];
672            let layer_type = if let Some(lsct) = layer.get_info(SECTION_DIVIDER_SETTING_KEY) {
673                let type_info = SectionDividerSetting::unpack(
674                    &mut MemReaderRef::new(lsct),
675                    true,
676                    psd.encoding,
677                    &None,
678                )?;
679                type_info.typ
680            } else {
681                0
682            };
683            if layer_type == 1 || layer_type == 2 {
684                let is_closed = layer_type == 2;
685                return Ok(GroupLayer {
686                    layer,
687                    layer_idx: *layer_idx,
688                    psd,
689                    childrens,
690                    is_closed,
691                });
692            } else if layer_type == 3 {
693                childrens.push(Layer::Group(GroupLayer::create(layer_idx, psd)?));
694            } else if layer_type == 0 {
695                childrens.push(Layer::Normal(NormalLayer {
696                    layer,
697                    layer_idx: *layer_idx,
698                    psd,
699                }));
700            } else {
701                anyhow::bail!("Unknown layer section divider type: {}", layer_type);
702            }
703            *layer_idx += 1;
704        }
705        anyhow::bail!("Layer group does not have a start marker");
706    }
707
708    /// Returns whether the layer group is closed.
709    pub fn is_closed(&self) -> bool {
710        self.is_closed
711    }
712
713    /// Returns the name of the layer group.
714    pub fn layer_name(&self) -> Result<String> {
715        self.layer.layer_name(self.psd.encoding)
716    }
717
718    /// Returns the current layer's index in the PSD file.
719    pub fn layer_index(&self) -> usize {
720        self.layer_idx
721    }
722}
723
724impl<'a> Deref for GroupLayer<'a> {
725    type Target = Vec<Layer<'a>>;
726
727    fn deref(&self) -> &Self::Target {
728        &self.childrens
729    }
730}
731
732#[derive(Debug)]
733pub enum Layer<'a> {
734    Normal(NormalLayer<'a>),
735    Group(GroupLayer<'a>),
736}
737
738impl<'a> Layer<'a> {
739    /// Returns the name of the layer.
740    pub fn layer_name(&self) -> Result<String> {
741        match self {
742            Layer::Normal(n) => n.layer_name(),
743            Layer::Group(g) => g.layer_name(),
744        }
745    }
746
747    /// Returns the current layer's index in the PSD file.
748    pub fn layer_index(&self) -> usize {
749        match self {
750            Layer::Normal(n) => n.layer_index(),
751            Layer::Group(g) => g.layer_index(),
752        }
753    }
754}
755
756impl PsdReader {
757    pub fn new<T: Read + Seek>(mut reader: T, encoding: Encoding) -> Result<Self> {
758        let psd = PsdFile::unpack(&mut reader, true, encoding, &None)?;
759        if psd.header.signature != *PSD_SIGNATURE {
760            anyhow::bail!("Invalid PSD signature");
761        }
762        if psd.header.version != 1 {
763            anyhow::bail!("Unsupported PSD version: {}", psd.header.version);
764        }
765        let mut channel_start_indices = Vec::new();
766        let mut idx = 0;
767        for layer in &psd.layer_and_mask_info.layer_info.layer_records {
768            channel_start_indices.push(idx);
769            idx += layer.base.channels as usize;
770        }
771        Ok(Self {
772            psd,
773            encoding,
774            channel_start_indices,
775        })
776    }
777
778    /// Returns the width of the PSD image.
779    pub fn width(&self) -> u32 {
780        self.psd.header.width
781    }
782
783    /// Returns the height of the PSD image.
784    pub fn height(&self) -> u32 {
785        self.psd.header.height
786    }
787
788    /// Returns the color mode of the PSD file.
789    ///
790    /// 1 = Grayscale, 3 = RGB, etc.
791    pub fn color_mode(&self) -> u16 {
792        self.psd.header.color_mode
793    }
794
795    /// Returns the number of channels in the PSD file.
796    pub fn channels(&self) -> u16 {
797        self.psd.header.channels
798    }
799
800    /// Returns the bit depth of the PSD file.
801    pub fn bit_depth(&self) -> u16 {
802        self.psd.header.depth
803    }
804
805    /// Reads and returns the layers in the PSD file.
806    pub fn read_layers<'a>(&'a self) -> Result<Vec<Layer<'a>>> {
807        let mut layers = Vec::new();
808        let mut layer_idx = 0;
809        let count = self.psd.layer_count();
810        while layer_idx < count {
811            let layer = &self.psd.layer_and_mask_info.layer_info.layer_records[layer_idx];
812            let layer_type = if let Some(lsct) = layer.get_info(SECTION_DIVIDER_SETTING_KEY) {
813                let type_info = SectionDividerSetting::unpack(
814                    &mut MemReaderRef::new(lsct),
815                    true,
816                    self.encoding,
817                    &None,
818                )?;
819                type_info.typ
820            } else {
821                0
822            };
823            if layer_type == 1 || layer_type == 2 {
824                anyhow::bail!("Layer group does not have an end marker");
825            }
826            if layer_type == 0 {
827                layers.push(Layer::Normal(NormalLayer {
828                    layer,
829                    layer_idx,
830                    psd: self,
831                }));
832            } else if layer_type == 3 {
833                layers.push(Layer::Group(GroupLayer::create(&mut layer_idx, self)?));
834            } else {
835                anyhow::bail!("Unknown layer section divider type: {}", layer_type);
836            }
837            layer_idx += 1;
838        }
839        Ok(layers)
840    }
841
842    /// Returns the normal layers in the PSD file. This function ignores layer groups.
843    pub fn read_normal_layers<'a>(&'a self) -> Result<Vec<NormalLayer<'a>>> {
844        let mut layers = Vec::new();
845        let count = self.psd.layer_count();
846        for layer_idx in 0..count {
847            let layer = &self.psd.layer_and_mask_info.layer_info.layer_records[layer_idx];
848            let layer_type = if let Some(lsct) = layer.get_info(SECTION_DIVIDER_SETTING_KEY) {
849                let type_info = SectionDividerSetting::unpack(
850                    &mut MemReaderRef::new(lsct),
851                    true,
852                    self.encoding,
853                    &None,
854                )?;
855                type_info.typ
856            } else {
857                0
858            };
859            if layer_type == 0 {
860                layers.push(NormalLayer {
861                    layer,
862                    layer_idx,
863                    psd: self,
864                });
865            }
866        }
867        Ok(layers)
868    }
869}
msg_tool\utils\psd/mod.rs

msg_tool\utils\psd/
mod.rs
