Skip to main content

aster/media/
image.rs

1//! 图片处理模块
2//!
3
4use base64::{engine::general_purpose::STANDARD, Engine};
5use std::collections::HashSet;
6use std::fs;
7use std::path::Path;
8use std::sync::LazyLock;
9
10use super::mime::get_mime_type_sync;
11
12/// 支持的图片格式
13pub static SUPPORTED_IMAGE_FORMATS: LazyLock<HashSet<&'static str>> =
14    LazyLock::new(|| HashSet::from(["png", "jpg", "jpeg", "gif", "webp"]));
15
16/// 最大图片 token 数
17pub const MAX_IMAGE_TOKENS: u64 = 25000;
18
19/// 图片压缩配置
20pub struct ImageCompressionConfig {
21    pub max_width: u32,
22    pub max_height: u32,
23    pub quality: u8,
24}
25
26pub const IMAGE_COMPRESSION_CONFIG: ImageCompressionConfig = ImageCompressionConfig {
27    max_width: 400,
28    max_height: 400,
29    quality: 20,
30};
31
32/// 图片尺寸信息
33#[derive(Debug, Clone, Default)]
34pub struct ImageDimensions {
35    pub original_width: Option<u32>,
36    pub original_height: Option<u32>,
37    pub display_width: Option<u32>,
38    pub display_height: Option<u32>,
39}
40
41/// 图片处理结果
42#[derive(Debug, Clone)]
43pub struct ImageResult {
44    pub base64: String,
45    pub mime_type: String,
46    pub original_size: u64,
47    pub dimensions: Option<ImageDimensions>,
48}
49
50/// 检查是否为支持的图片格式
51pub fn is_supported_image_format(ext: &str) -> bool {
52    let normalized = ext.to_lowercase().replace('.', "");
53    SUPPORTED_IMAGE_FORMATS.contains(normalized.as_str())
54}
55
56/// 估算图片的 token 消耗
57pub fn estimate_image_tokens(base64: &str) -> u64 {
58    (base64.len() as f64 * 0.125).ceil() as u64
59}
60
61/// 读取图片文件(同步版本,不压缩)
62pub fn read_image_file_sync(file_path: &Path) -> Result<ImageResult, String> {
63    let metadata =
64        fs::metadata(file_path).map_err(|e| format!("Failed to read file metadata: {}", e))?;
65
66    if metadata.len() == 0 {
67        return Err(format!("Image file is empty: {}", file_path.display()));
68    }
69
70    let buffer = fs::read(file_path).map_err(|e| format!("Failed to read file: {}", e))?;
71
72    let ext = file_path
73        .extension()
74        .and_then(|e| e.to_str())
75        .unwrap_or("png")
76        .to_lowercase();
77
78    let mime_type = get_mime_type_sync(&buffer)
79        .unwrap_or_else(|| Box::leak(format!("image/{}", ext).into_boxed_str()));
80
81    let base64 = STANDARD.encode(&buffer);
82
83    Ok(ImageResult {
84        base64,
85        mime_type: mime_type.to_string(),
86        original_size: metadata.len(),
87        dimensions: None,
88    })
89}
90
91/// 验证图片文件
92pub fn validate_image_file(file_path: &Path) -> Result<(), String> {
93    if !file_path.exists() {
94        return Err("File does not exist".to_string());
95    }
96
97    let metadata =
98        fs::metadata(file_path).map_err(|e| format!("Failed to read metadata: {}", e))?;
99
100    if metadata.len() == 0 {
101        return Err("Image file is empty".to_string());
102    }
103
104    let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or("");
105
106    if !is_supported_image_format(ext) {
107        return Err(format!(
108            "Unsupported image format: {}. Supported: {:?}",
109            ext,
110            SUPPORTED_IMAGE_FORMATS.iter().collect::<Vec<_>>()
111        ));
112    }
113
114    Ok(())
115}
116
117/// 读取图片文件(增强版本,包含尺寸提取)
118///
119/// 参考 claude-code-open 实现,提供更详细的图片信息
120pub fn read_image_file_enhanced(file_path: &Path) -> Result<ImageResult, String> {
121    let metadata =
122        fs::metadata(file_path).map_err(|e| format!("Failed to read file metadata: {}", e))?;
123
124    if metadata.len() == 0 {
125        return Err(format!("Image file is empty: {}", file_path.display()));
126    }
127
128    let buffer = fs::read(file_path).map_err(|e| format!("Failed to read file: {}", e))?;
129
130    let ext = file_path
131        .extension()
132        .and_then(|e| e.to_str())
133        .unwrap_or("png")
134        .to_lowercase();
135
136    let mime_type = get_mime_type_sync(&buffer)
137        .unwrap_or_else(|| Box::leak(format!("image/{}", ext).into_boxed_str()));
138
139    let base64 = STANDARD.encode(&buffer);
140
141    // 计算 token 估算
142    let _token_estimate = estimate_image_tokens(&base64);
143
144    // 尝试提取图片尺寸(基于文件大小和格式)
145    let dimensions = estimate_image_dimensions(&buffer, metadata.len());
146
147    Ok(ImageResult {
148        base64,
149        mime_type: mime_type.to_string(),
150        original_size: metadata.len(),
151        dimensions: Some(dimensions),
152    })
153}
154
155/// 估算图片尺寸(基于文件大小和格式)
156///
157/// 这是一个简化版本,不依赖外部图像处理库
158/// 实际项目中可以添加 image-rs 或 sharp 等库进行精确提取
159pub fn estimate_image_dimensions(_buffer: &[u8], file_size: u64) -> ImageDimensions {
160    // TODO: 集成 image-rs 或 sharp 库来提取实际尺寸
161    //
162    // 需要在 Cargo.toml 中添加:
163    // image-rs = { version = "0.25", features = ["jpeg", "png", "gif", "webp"] }
164    //
165    // 然后使用:
166    // let reader = image::ImageReader::new(Cursor::new(buffer))
167    //     .with_guessed_format();
168    // let dimensions = reader.dimensions().unwrap();
169    //
170    // 暂时基于文件大小估算(非常粗略)
171    let estimated_pixels = file_size / 3; // 假设每个像素平均 3 字节(RGB)
172    let estimated_size = (estimated_pixels as f64).sqrt() as u32;
173    let size = estimated_size.max(100); // 最小 100x100
174
175    ImageDimensions {
176        original_width: Some(size),
177        original_height: Some(size),
178        display_width: Some(size),
179        display_height: Some(size),
180    }
181}