zerolaunch_plugin_api/common/
image_utils.rs

1use image::GenericImageView;
2use image::ImageBuffer;
3use image::ImageEncoder;
4use image::ImageFormat;
5use image::Rgba;
6use image::RgbaImage;
7use kmeans_colors::get_kmeans;
8use palette::{IntoColor, Lab, Srgb};
9use rand::RngExt;
10use rayon::prelude::*;
11use std::io::Cursor;
12use std::sync::{Arc, Mutex};
13use tracing::info;
14
15/// 跨平台图片处理工具的错误类型。
16/// 不依赖 crate::error，SDK 内部自足。
17#[derive(Debug, thiserror::Error)]
18pub enum ImageUtilsError {
19    #[error("图片处理错误: {0}")]
20    ProcessingError(String),
21
22    #[error("任务执行错误: {0}")]
23    TaskJoinError(String),
24}
25
26/// 跨平台图片处理工具函数集合。
27/// 从 core::image_processor 中提取的平台无关逻辑，不依赖任何 Windows API。
28pub struct ImageUtils;
29
30impl ImageUtils {
31    /// 将图片数据转换为 PNG 格式。
32    /// 支持 SVG、BMP、JPEG、GIF、WebP 等格式自动检测并转换。
33    /// 参数：image_data - 原始图片字节数据。
34    /// 返回：PNG 格式字节数据，失败返回 ImageUtilsError。
35    pub async fn convert_image_to_png(image_data: Vec<u8>) -> Result<Vec<u8>, ImageUtilsError> {
36        if image_data.is_empty() {
37            return Err(ImageUtilsError::ProcessingError(
38                "Input image data is empty".to_string(),
39            ));
40        }
41
42        if Self::is_html_content(&image_data) {
43            return Err(ImageUtilsError::ProcessingError(
44                "Downloaded content appears to be HTML, not an image".to_string(),
45            ));
46        }
47
48        tokio::task::spawn_blocking(move || -> Result<Vec<u8>, ImageUtilsError> {
49            match usvg::Tree::from_data(&image_data, &usvg::Options::default()) {
50                Ok(tree) => {
51                    let pixmap_size = tree.size().to_int_size();
52                    if pixmap_size.width() == 0 || pixmap_size.height() == 0 {
53                        return Err(ImageUtilsError::ProcessingError(format!(
54                            "Invalid SVG dimensions (width: {}px, height: {}px)",
55                            pixmap_size.width(),
56                            pixmap_size.height()
57                        )));
58                    }
59
60                    let mut pixmap =
61                        tiny_skia::Pixmap::new(pixmap_size.width(), pixmap_size.height())
62                            .ok_or_else(|| {
63                                ImageUtilsError::ProcessingError(
64                                    "Failed to create Pixmap for SVG rendering".to_string(),
65                                )
66                            })?;
67
68                    resvg::render(&tree, tiny_skia::Transform::default(), &mut pixmap.as_mut());
69
70                    let png_data = pixmap.encode_png().map_err(|e| {
71                        ImageUtilsError::ProcessingError(format!(
72                            "Failed to encode SVG as PNG: {}",
73                            e
74                        ))
75                    })?;
76                    Ok(png_data)
77                }
78                Err(_) => {
79                    let img_reader = image::ImageReader::new(Cursor::new(image_data))
80                        .with_guessed_format()
81                        .map_err(|e| {
82                            ImageUtilsError::ProcessingError(format!(
83                                "Failed to create image reader: {}",
84                                e
85                            ))
86                        })?;
87
88                    let format = img_reader.format().ok_or_else(|| {
89                        ImageUtilsError::ProcessingError(
90                            "Unable to detect image format".to_string(),
91                        )
92                    })?;
93
94                    let mut img = img_reader.decode().map_err(|e| {
95                        ImageUtilsError::ProcessingError(format!("Failed to decode image: {}", e))
96                    })?;
97
98                    if format != ImageFormat::Png {
99                        img = image::DynamicImage::ImageRgba8(img.to_rgba8());
100                    }
101
102                    let mut png_data = Vec::new();
103                    let encoder = image::codecs::png::PngEncoder::new(&mut png_data);
104                    img.write_with_encoder(encoder).map_err(|e| {
105                        ImageUtilsError::ProcessingError(format!(
106                            "Failed to encode image as PNG: {}",
107                            e
108                        ))
109                    })?;
110                    Ok(png_data)
111                }
112            }
113        })
114        .await
115        .map_err(|e| ImageUtilsError::TaskJoinError(format!("Task join error: {}", e)))?
116    }
117
118    /// 将 PNG 图片数据编码为 WebP 格式（无损 VP8L）。
119    /// 用 image crate 内置的纯 Rust 无损编码器（image-webp），不依赖 libwebp C 库。
120    /// 参数：png_data - PNG 格式图片字节数据。
121    /// 返回：WebP 格式字节数据，失败返回 ImageUtilsError。
122    pub fn to_webp(png_data: Vec<u8>) -> Result<Vec<u8>, ImageUtilsError> {
123        let img = image::load_from_memory(&png_data).map_err(|e| {
124            ImageUtilsError::ProcessingError(format!("Failed to load image for webp: {}", e))
125        })?;
126
127        let mut webp_data = Vec::new();
128        let encoder = image::codecs::webp::WebPEncoder::new_lossless(&mut webp_data);
129        img.write_with_encoder(encoder).map_err(|e| {
130            ImageUtilsError::ProcessingError(format!("Failed to encode image as WebP: {}", e))
131        })?;
132        Ok(webp_data)
133    }
134
135    /// 将图片字节数据转换为 base64 data URL（MIME 由 image crate 的字节头嗅探决定）。
136    /// 参数：data - 图片字节数据。
137    /// 返回：data URL 字符串；空数据返回空字符串；无法识别格式时按 PNG 标注。
138    pub fn to_data_url(data: &[u8]) -> String {
139        if data.is_empty() {
140            return String::new();
141        }
142        use base64::Engine;
143        let mime = image::guess_format(data)
144            .map(|format| format.to_mime_type())
145            .unwrap_or("image/png");
146        format!(
147            "data:{};base64,{}",
148            mime,
149            base64::engine::general_purpose::STANDARD.encode(data)
150        )
151    }
152
153    /// 判断数据是否像是 HTML 内容。
154    /// 参数：data - 待检测的字节数据。
155    /// 返回：true 表示是 HTML 内容。
156    pub fn is_html_content(data: &[u8]) -> bool {
157        if let Ok(s) = std::str::from_utf8(data) {
158            let s_trimmed = s.trim_start();
159            if s_trimmed.eq_ignore_ascii_case("<!DOCTYPE html")
160                || s_trimmed.starts_with("<!DOCTYPE html")
161                || s_trimmed.starts_with("<html")
162                || s_trimmed.starts_with("<HTML")
163            {
164                return true;
165            }
166        }
167        false
168    }
169
170    /// 调整图片大小，如果超过指定尺寸则等比缩放。
171    /// 参数：data - PNG 图片字节数据；max_width - 最大宽度；max_height - 最大高度。
172    /// 返回：调整大小后的 PNG 字节数据，失败返回 ImageUtilsError。
173    pub async fn resize_image(
174        data: Vec<u8>,
175        max_width: u32,
176        max_height: u32,
177    ) -> Result<Vec<u8>, ImageUtilsError> {
178        tokio::task::spawn_blocking(move || -> Result<Vec<u8>, ImageUtilsError> {
179            let img = image::load_from_memory(&data).map_err(|e| {
180                ImageUtilsError::ProcessingError(format!(
181                    "Failed to load image for resizing: {}",
182                    e
183                ))
184            })?;
185
186            if img.width() <= max_width && img.height() <= max_height {
187                return Ok(data);
188            }
189
190            let resized = img.thumbnail(max_width, max_height);
191
192            let mut png_data = Vec::new();
193            let encoder = image::codecs::png::PngEncoder::new(&mut png_data);
194            resized.write_with_encoder(encoder).map_err(|e| {
195                ImageUtilsError::ProcessingError(format!("Failed to encode resized image: {}", e))
196            })?;
197
198            Ok(png_data)
199        })
200        .await
201        .map_err(|e| ImageUtilsError::TaskJoinError(format!("Task join error: {}", e)))?
202    }
203
204    /// 从 PNG 图像数据中裁剪掉外围的白色或透明像素。
205    /// 参数：png_data - PNG 格式图片字节数据（须为正方形）。
206    /// 返回：裁剪后的 PNG 字节数据，失败返回 ImageUtilsError。
207    pub fn trim_transparent_white_border(png_data: Vec<u8>) -> Result<Vec<u8>, ImageUtilsError> {
208        let img = image::load_from_memory(&png_data).map_err(|e| {
209            ImageUtilsError::ProcessingError(format!("Failed to load image from memory: {}", e))
210        })?;
211
212        let width = img.width();
213        let height = img.height();
214
215        if width != height {
216            return Err(ImageUtilsError::ProcessingError(format!(
217                "Input image is not square: {}x{}",
218                width, height
219            )));
220        }
221
222        let mut border_width = 0;
223        let size = width;
224
225        'outer: for layer in 0..size / 2 {
226            for x in layer..size - layer {
227                let pixel = img.get_pixel(x, layer);
228                if !Self::is_white_or_transparent(pixel) {
229                    break 'outer;
230                }
231            }
232
233            for y in layer..size - layer {
234                let pixel = img.get_pixel(size - 1 - layer, y);
235                if !Self::is_white_or_transparent(pixel) {
236                    break 'outer;
237                }
238            }
239
240            for x in layer..size - layer {
241                let pixel = img.get_pixel(x, size - 1 - layer);
242                if !Self::is_white_or_transparent(pixel) {
243                    break 'outer;
244                }
245            }
246
247            for y in layer..size - layer {
248                let pixel = img.get_pixel(layer, y);
249                if !Self::is_white_or_transparent(pixel) {
250                    break 'outer;
251                }
252            }
253
254            border_width = layer + 1;
255        }
256
257        if border_width >= size / 2 {
258            return Ok(png_data);
259        }
260
261        let new_size = size - 2 * border_width;
262        let mut new_img = ImageBuffer::new(new_size, new_size);
263
264        for y in 0..new_size {
265            for x in 0..new_size {
266                let pixel = img.get_pixel(x + border_width, y + border_width);
267                new_img.put_pixel(x, y, pixel);
268            }
269        }
270
271        let mut output = Vec::new();
272        let encoder = image::codecs::png::PngEncoder::new(&mut output);
273        encoder
274            .write_image(
275                &new_img.into_raw(),
276                new_size,
277                new_size,
278                image::ColorType::Rgba8.into(),
279            )
280            .map_err(|e| {
281                ImageUtilsError::ProcessingError(format!("Failed to encode trimmed image: {}", e))
282            })?;
283
284        Ok(output)
285    }
286
287    /// 判断像素是否为白色或透明。
288    fn is_white_or_transparent(pixel: Rgba<u8>) -> bool {
289        pixel[3] < 10 || (pixel[0] > 245 && pixel[1] > 245 && pixel[2] > 245)
290    }
291
292    /// 将 RGBA 图像数据编码为 PNG 格式。
293    /// 参数：rgba_image - RGBA 格式图像缓冲区。
294    /// 返回：PNG 字节数据，失败返回 ImageUtilsError。
295    pub fn rgba_image_to_png(rgba_image: &RgbaImage) -> Result<Vec<u8>, ImageUtilsError> {
296        let mut buffer = Vec::new();
297        let mut cursor = Cursor::new(&mut buffer);
298        rgba_image
299            .write_to(&mut cursor, ImageFormat::Png)
300            .map_err(|e| {
301                ImageUtilsError::ProcessingError(format!(
302                    "Failed to encode RGBA image as PNG: {}",
303                    e
304                ))
305            })?;
306        Ok(buffer)
307    }
308
309    /// 获取图片的主色调。
310    /// 使用 K-Means 聚类算法分析图片中非透明像素的颜色分布。
311    /// 参数：image_data - PNG 格式图片字节数据。
312    /// 返回：(R, G, B) 主色调元组，失败返回 ImageUtilsError。
313    pub async fn get_dominant_color(image_data: Vec<u8>) -> Result<(u8, u8, u8), ImageUtilsError> {
314        tokio::task::spawn_blocking(move || -> Result<(u8, u8, u8), ImageUtilsError> {
315            let img = image::load_from_memory(&image_data).map_err(|e| {
316                ImageUtilsError::ProcessingError(format!(
317                    "Failed to load image for color analysis: {}",
318                    e
319                ))
320            })?;
321            let rgba_img = img.to_rgba8();
322
323            let pixels: Vec<[u8; 3]> = rgba_img
324                .pixels()
325                .par_bridge()
326                .filter_map(|pixel| {
327                    if pixel[3] != 0 {
328                        Some([pixel[0], pixel[1], pixel[2]])
329                    } else {
330                        None
331                    }
332                })
333                .collect();
334
335            if pixels.is_empty() {
336                return Err(ImageUtilsError::ProcessingError(
337                    "No visible pixels found in image for color analysis".to_string(),
338                ));
339            }
340
341            let lab_samples: Vec<Lab> = pixels
342                .par_iter()
343                .map(|&rgb| Srgb::from(rgb).into_format::<f32>().into_color())
344                .collect();
345
346            let cluster_count = 5;
347            let max_iterations = 20;
348            let tolerance = 1.0;
349            let runs = 3;
350            let verbose = false;
351
352            let lab_samples_arc = Arc::new(lab_samples);
353            let best_result = (0..runs)
354                .into_par_iter()
355                .map(|_| {
356                    let seed = rand::rng().random::<u64>();
357                    let samples = Arc::clone(&lab_samples_arc);
358                    get_kmeans(
359                        cluster_count,
360                        max_iterations,
361                        tolerance,
362                        verbose,
363                        &samples,
364                        seed,
365                    )
366                })
367                .min_by(|a, b| {
368                    a.score
369                        .partial_cmp(&b.score)
370                        .expect("Score comparison should not fail")
371                })
372                .expect("Should have at least one clustering result");
373
374            let cluster_counts = {
375                let counts = vec![0; cluster_count];
376                let counts_mutex = Arc::new(Mutex::new(counts));
377
378                best_result.indices.par_iter().for_each(|&i| {
379                    let mut counts = counts_mutex.lock().expect("Mutex should not be poisoned");
380                    counts[i as usize] += 1;
381                });
382
383                Arc::try_unwrap(counts_mutex)
384                    .expect("Arc should have only one reference")
385                    .into_inner()
386                    .expect("Mutex should not be poisoned")
387            };
388
389            let (dominant_idx, _) = cluster_counts
390                .iter()
391                .enumerate()
392                .max_by_key(|&(_, count)| count)
393                .expect("Should have at least one cluster count");
394
395            let dominant_lab = best_result.centroids[dominant_idx];
396            let srgb: Srgb = dominant_lab.into_color();
397            let rgb = srgb.into_format::<u8>();
398            let (r, g, b) = rgb.into_components();
399
400            info!("Dominant color analysis complete: RGB({}, {}, {})", r, g, b);
401            Ok((r, g, b))
402        })
403        .await
404        .map_err(|e| ImageUtilsError::TaskJoinError(format!("Task join error: {}", e)))?
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use super::ImageUtils;
411
412    /// 各格式的字节头（image crate 的 guess_format 只读魔数，不需要完整文件）。
413    #[test]
414    fn to_data_url_labels_bytes_by_guessed_format() {
415        let cases: [(&[u8], &str); 5] = [
416            (b"RIFF\x00\x00\x00\x00WEBPVP8L", "data:image/webp;base64,"),
417            (
418                b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR",
419                "data:image/png;base64,",
420            ),
421            (
422                &[0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10],
423                "data:image/jpeg;base64,",
424            ),
425            (b"GIF89a\x01\x00", "data:image/gif;base64,"),
426            (
427                &[0x00, 0x00, 0x01, 0x00, 0x01, 0x00],
428                "data:image/x-icon;base64,",
429            ),
430        ];
431        for (bytes, prefix) in cases {
432            let url = ImageUtils::to_data_url(bytes);
433            assert!(
434                url.starts_with(prefix),
435                "字节头 {bytes:?} 期望 {prefix}，实际 {url}"
436            );
437        }
438    }
439
440    #[test]
441    fn to_data_url_falls_back_to_png_for_unknown_bytes() {
442        assert!(ImageUtils::to_data_url(b"not an image").starts_with("data:image/png;base64,"));
443        assert_eq!(ImageUtils::to_data_url(b""), "");
444    }
445}
zerolaunch_plugin_api/common/image_utils.rs

zerolaunch_plugin_api/common/
image_utils.rs
