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 前缀。
136    /// WebP 字节头为 "RIFF....WEBP"，否则按 PNG 处理。
137    /// 参数：data - 图片字节数据。
138    /// 返回：data URL 前缀（"data:image/webp;base64," / "data:image/png;base64,"），未知回退 PNG。
139    pub fn data_url_prefix(data: &[u8]) -> &'static str {
140        if data.len() >= 12 && &data[0..4] == b"RIFF" && &data[8..12] == b"WEBP" {
141            "data:image/webp;base64,"
142        } else {
143            "data:image/png;base64,"
144        }
145    }
146
147    /// 将图片字节数据转换为 base64 data URL（MIME 按字节头嗅探）。
148    /// 参数：data - 图片字节数据。
149    /// 返回：data URL 字符串；空数据返回空字符串。
150    pub fn to_data_url(data: &[u8]) -> String {
151        if data.is_empty() {
152            return String::new();
153        }
154        use base64::Engine;
155        format!(
156            "{}{}",
157            Self::data_url_prefix(data),
158            base64::engine::general_purpose::STANDARD.encode(data)
159        )
160    }
161
162    /// 判断数据是否像是 HTML 内容。
163    /// 参数：data - 待检测的字节数据。
164    /// 返回：true 表示是 HTML 内容。
165    pub fn is_html_content(data: &[u8]) -> bool {
166        if let Ok(s) = std::str::from_utf8(data) {
167            let s_trimmed = s.trim_start();
168            if s_trimmed.eq_ignore_ascii_case("<!DOCTYPE html")
169                || s_trimmed.starts_with("<!DOCTYPE html")
170                || s_trimmed.starts_with("<html")
171                || s_trimmed.starts_with("<HTML")
172            {
173                return true;
174            }
175        }
176        false
177    }
178
179    /// 调整图片大小，如果超过指定尺寸则等比缩放。
180    /// 参数：data - PNG 图片字节数据；max_width - 最大宽度；max_height - 最大高度。
181    /// 返回：调整大小后的 PNG 字节数据，失败返回 ImageUtilsError。
182    pub async fn resize_image(
183        data: Vec<u8>,
184        max_width: u32,
185        max_height: u32,
186    ) -> Result<Vec<u8>, ImageUtilsError> {
187        tokio::task::spawn_blocking(move || -> Result<Vec<u8>, ImageUtilsError> {
188            let img = image::load_from_memory(&data).map_err(|e| {
189                ImageUtilsError::ProcessingError(format!(
190                    "Failed to load image for resizing: {}",
191                    e
192                ))
193            })?;
194
195            if img.width() <= max_width && img.height() <= max_height {
196                return Ok(data);
197            }
198
199            let resized = img.thumbnail(max_width, max_height);
200
201            let mut png_data = Vec::new();
202            let encoder = image::codecs::png::PngEncoder::new(&mut png_data);
203            resized.write_with_encoder(encoder).map_err(|e| {
204                ImageUtilsError::ProcessingError(format!("Failed to encode resized image: {}", e))
205            })?;
206
207            Ok(png_data)
208        })
209        .await
210        .map_err(|e| ImageUtilsError::TaskJoinError(format!("Task join error: {}", e)))?
211    }
212
213    /// 从 PNG 图像数据中裁剪掉外围的白色或透明像素。
214    /// 参数：png_data - PNG 格式图片字节数据（须为正方形）。
215    /// 返回：裁剪后的 PNG 字节数据，失败返回 ImageUtilsError。
216    pub fn trim_transparent_white_border(png_data: Vec<u8>) -> Result<Vec<u8>, ImageUtilsError> {
217        let img = image::load_from_memory(&png_data).map_err(|e| {
218            ImageUtilsError::ProcessingError(format!("Failed to load image from memory: {}", e))
219        })?;
220
221        let width = img.width();
222        let height = img.height();
223
224        if width != height {
225            return Err(ImageUtilsError::ProcessingError(format!(
226                "Input image is not square: {}x{}",
227                width, height
228            )));
229        }
230
231        let mut border_width = 0;
232        let size = width;
233
234        'outer: for layer in 0..size / 2 {
235            for x in layer..size - layer {
236                let pixel = img.get_pixel(x, layer);
237                if !Self::is_white_or_transparent(pixel) {
238                    break 'outer;
239                }
240            }
241
242            for y in layer..size - layer {
243                let pixel = img.get_pixel(size - 1 - layer, y);
244                if !Self::is_white_or_transparent(pixel) {
245                    break 'outer;
246                }
247            }
248
249            for x in layer..size - layer {
250                let pixel = img.get_pixel(x, size - 1 - layer);
251                if !Self::is_white_or_transparent(pixel) {
252                    break 'outer;
253                }
254            }
255
256            for y in layer..size - layer {
257                let pixel = img.get_pixel(layer, y);
258                if !Self::is_white_or_transparent(pixel) {
259                    break 'outer;
260                }
261            }
262
263            border_width = layer + 1;
264        }
265
266        if border_width >= size / 2 {
267            return Ok(png_data);
268        }
269
270        let new_size = size - 2 * border_width;
271        let mut new_img = ImageBuffer::new(new_size, new_size);
272
273        for y in 0..new_size {
274            for x in 0..new_size {
275                let pixel = img.get_pixel(x + border_width, y + border_width);
276                new_img.put_pixel(x, y, pixel);
277            }
278        }
279
280        let mut output = Vec::new();
281        let encoder = image::codecs::png::PngEncoder::new(&mut output);
282        encoder
283            .write_image(
284                &new_img.into_raw(),
285                new_size,
286                new_size,
287                image::ColorType::Rgba8.into(),
288            )
289            .map_err(|e| {
290                ImageUtilsError::ProcessingError(format!("Failed to encode trimmed image: {}", e))
291            })?;
292
293        Ok(output)
294    }
295
296    /// 判断像素是否为白色或透明。
297    fn is_white_or_transparent(pixel: Rgba<u8>) -> bool {
298        pixel[3] < 10 || (pixel[0] > 245 && pixel[1] > 245 && pixel[2] > 245)
299    }
300
301    /// 将 RGBA 图像数据编码为 PNG 格式。
302    /// 参数：rgba_image - RGBA 格式图像缓冲区。
303    /// 返回：PNG 字节数据，失败返回 ImageUtilsError。
304    pub fn rgba_image_to_png(rgba_image: &RgbaImage) -> Result<Vec<u8>, ImageUtilsError> {
305        let mut buffer = Vec::new();
306        let mut cursor = Cursor::new(&mut buffer);
307        rgba_image
308            .write_to(&mut cursor, ImageFormat::Png)
309            .map_err(|e| {
310                ImageUtilsError::ProcessingError(format!(
311                    "Failed to encode RGBA image as PNG: {}",
312                    e
313                ))
314            })?;
315        Ok(buffer)
316    }
317
318    /// 获取图片的主色调。
319    /// 使用 K-Means 聚类算法分析图片中非透明像素的颜色分布。
320    /// 参数：image_data - PNG 格式图片字节数据。
321    /// 返回：(R, G, B) 主色调元组，失败返回 ImageUtilsError。
322    pub async fn get_dominant_color(image_data: Vec<u8>) -> Result<(u8, u8, u8), ImageUtilsError> {
323        tokio::task::spawn_blocking(move || -> Result<(u8, u8, u8), ImageUtilsError> {
324            let img = image::load_from_memory(&image_data).map_err(|e| {
325                ImageUtilsError::ProcessingError(format!(
326                    "Failed to load image for color analysis: {}",
327                    e
328                ))
329            })?;
330            let rgba_img = img.to_rgba8();
331
332            let pixels: Vec<[u8; 3]> = rgba_img
333                .pixels()
334                .par_bridge()
335                .filter_map(|pixel| {
336                    if pixel[3] != 0 {
337                        Some([pixel[0], pixel[1], pixel[2]])
338                    } else {
339                        None
340                    }
341                })
342                .collect();
343
344            if pixels.is_empty() {
345                return Err(ImageUtilsError::ProcessingError(
346                    "No visible pixels found in image for color analysis".to_string(),
347                ));
348            }
349
350            let lab_samples: Vec<Lab> = pixels
351                .par_iter()
352                .map(|&rgb| Srgb::from(rgb).into_format::<f32>().into_color())
353                .collect();
354
355            let cluster_count = 5;
356            let max_iterations = 20;
357            let tolerance = 1.0;
358            let runs = 3;
359            let verbose = false;
360
361            let lab_samples_arc = Arc::new(lab_samples);
362            let best_result = (0..runs)
363                .into_par_iter()
364                .map(|_| {
365                    let seed = rand::rng().random::<u64>();
366                    let samples = Arc::clone(&lab_samples_arc);
367                    get_kmeans(
368                        cluster_count,
369                        max_iterations,
370                        tolerance,
371                        verbose,
372                        &samples,
373                        seed,
374                    )
375                })
376                .min_by(|a, b| {
377                    a.score
378                        .partial_cmp(&b.score)
379                        .expect("Score comparison should not fail")
380                })
381                .expect("Should have at least one clustering result");
382
383            let cluster_counts = {
384                let counts = vec![0; cluster_count];
385                let counts_mutex = Arc::new(Mutex::new(counts));
386
387                best_result.indices.par_iter().for_each(|&i| {
388                    let mut counts = counts_mutex.lock().expect("Mutex should not be poisoned");
389                    counts[i as usize] += 1;
390                });
391
392                Arc::try_unwrap(counts_mutex)
393                    .expect("Arc should have only one reference")
394                    .into_inner()
395                    .expect("Mutex should not be poisoned")
396            };
397
398            let (dominant_idx, _) = cluster_counts
399                .iter()
400                .enumerate()
401                .max_by_key(|&(_, count)| count)
402                .expect("Should have at least one cluster count");
403
404            let dominant_lab = best_result.centroids[dominant_idx];
405            let srgb: Srgb = dominant_lab.into_color();
406            let rgb = srgb.into_format::<u8>();
407            let (r, g, b) = rgb.into_components();
408
409            info!("Dominant color analysis complete: RGB({}, {}, {})", r, g, b);
410            Ok((r, g, b))
411        })
412        .await
413        .map_err(|e| ImageUtilsError::TaskJoinError(format!("Task join error: {}", e)))?
414    }
415}
zerolaunch_plugin_api/common/image_utils.rs

zerolaunch_plugin_api/common/
image_utils.rs
