|
| 1 | +//! 图文会话截图保存 + OCR 编排(v0.11.7,ADR-020)。 |
| 2 | +//! |
| 3 | +//! @ai-context: 用户框选截屏 → base64(PNG 无损,前端 canvas 裁剪)→ 解码验证 → |
| 4 | +//! 存图(full+thumb,双指纹去重)→ 全图 OCR(960px 口径,bbox 反算) |
| 5 | +//! → OCR 块落库(region=full,screen_id=None 走视图层聚类)。 |
| 6 | +//! @ai-context: 与实时/导入链路 OCR 口径统一(960px 缩放 + score≥0.5);与自动 |
| 7 | +//! 采帧不同:用户显式框选的内容不过 UI 垃圾黑名单(意图即内容)。 |
| 8 | +//! @ai-context: OCR 失败不阻断截图保存(图已落盘,block_count=0 诚实返回)。 |
| 9 | +
|
| 10 | +use base64::Engine; |
| 11 | +use image::GenericImageView; |
| 12 | +use serde::Serialize; |
| 13 | + |
| 14 | +use crate::db::Db; |
| 15 | +use crate::engine::EnginePool; |
| 16 | +use crate::image_store::SessionImageStore; |
| 17 | +use crate::types::{NewSessionOcrBlock, TextBox}; |
| 18 | + |
| 19 | +/// OCR 输入最大宽度(与实时/导入链路同口径——推理成本近平方下降)。 |
| 20 | +const OCR_MAX_WIDTH: u32 = 960; |
| 21 | +/// OCR 块最低置信度(与实时/导入链路同口径)。 |
| 22 | +const MIN_SCORE: f32 = 0.5; |
| 23 | +/// 识别预览条数上限(前端即时反馈)。 |
| 24 | +const PREVIEW_MAX: usize = 3; |
| 25 | +/// 单图 base64 解码后字节上限(16MB——解码前拒绝,省 CPU)。 |
| 26 | +const MAX_PHOTO_BYTES: usize = 16 * 1024 * 1024; |
| 27 | + |
| 28 | +/// 截图保存结果(前端反馈契约;camelCase 与前端 PhotoResult 对齐)。 |
| 29 | +#[derive(Debug, Clone, PartialEq, Serialize)] |
| 30 | +#[serde(rename_all = "camelCase")] |
| 31 | +pub struct PhotoCaptureResult { |
| 32 | + /// 双指纹命中(与最近保存图相同,未重复存图/未重复落块) |
| 33 | + pub duplicated: bool, |
| 34 | + /// 新落库 OCR 块数(OCR 失败/超时可能为 0) |
| 35 | + pub block_count: usize, |
| 36 | + /// 识别文本预览(≤3 条;空 = 未识别出文字) |
| 37 | + pub preview: Vec<String>, |
| 38 | + /// 图片相对路径(full/<ts>.webp) |
| 39 | + pub image_ref: String, |
| 40 | +} |
| 41 | + |
| 42 | +/// 保存截图并 OCR(主入口;注入 store/DB/引擎池——端到端可测)。 |
| 43 | +/// |
| 44 | +/// @ai-context: store 由调用方跨命令持有(命令层互斥槽内复用)——双指纹去重 |
| 45 | +/// FIFO 是内存态,每次 new 会清空导致去重失效(实时链路即长驻 |
| 46 | +/// store);budget 计数同样跨截图累计(50 张预算会话级生效)。 |
| 47 | +/// @ai-context: OCR 有界等待 ≤20s——调用方在锁内调用时锁时长与此相当; |
| 48 | +/// 图文采集为用户串行动线(同刻仅一个截图保存),可接受。 |
| 49 | +pub fn save_photo_capture( |
| 50 | + store: &mut SessionImageStore, |
| 51 | + db: &Db, |
| 52 | + engines: &EnginePool, |
| 53 | + session_id: i64, |
| 54 | + image_b64: &str, |
| 55 | + timestamp_ms: u64, |
| 56 | +) -> Result<PhotoCaptureResult, String> { |
| 57 | + let img = decode_image(image_b64)?; |
| 58 | + let (w, h) = img.dimensions(); |
| 59 | + // ① 存图:RGB → BGRA8(save_frame 输入口径)→ full+thumb + 双指纹去重 |
| 60 | + let bgra: Vec<u8> = img.pixels().flat_map(|p| [p[2], p[1], p[0], 255u8]).collect(); |
| 61 | + let rel = store |
| 62 | + .save_frame(timestamp_ms, &bgra, w, h) |
| 63 | + .map_err(|e| format!("保存截图失败: {}", e))?; |
| 64 | + // 重复判定:save_frame 对重复帧返回已有路径(旧时间戳)——文件名 ≠ 本次 ts 即重复 |
| 65 | + let duplicated = rel != format!("full/{}.webp", timestamp_ms); |
| 66 | + if duplicated { |
| 67 | + return Ok(PhotoCaptureResult { duplicated: true, block_count: 0, preview: Vec::new(), image_ref: rel }); |
| 68 | + } |
| 69 | + // ② 全图 OCR(960px 口径;失败/超时降级——图已保存) |
| 70 | + let (blocks, scale) = ocr_blocks(&img, engines); |
| 71 | + // ③ 落库(不过 UI 垃圾黑名单——用户显式框选的内容即意图) |
| 72 | + let mut count = 0usize; |
| 73 | + let mut preview = Vec::new(); |
| 74 | + for b in blocks { |
| 75 | + if b.score < MIN_SCORE || b.text.trim().is_empty() { |
| 76 | + continue; |
| 77 | + } |
| 78 | + count += 1; |
| 79 | + if preview.len() < PREVIEW_MAX { |
| 80 | + preview.push(b.text.clone()); |
| 81 | + } |
| 82 | + if let Err(e) = db.add_ocr_block(&NewSessionOcrBlock { |
| 83 | + session_id, |
| 84 | + timestamp_ms, |
| 85 | + text: b.text, |
| 86 | + score: b.score, |
| 87 | + region: "full".to_string(), |
| 88 | + region_kind: None, |
| 89 | + bbox: b.bbox.map(|bb| TextBox { |
| 90 | + x: bb.x * scale, |
| 91 | + y: bb.y * scale, |
| 92 | + w: bb.w * scale, |
| 93 | + h: bb.h * scale, |
| 94 | + }), |
| 95 | + screen_id: None, |
| 96 | + }) { |
| 97 | + eprintln!("[Photo] OCR 块落库失败: {}", e); |
| 98 | + } |
| 99 | + } |
| 100 | + Ok(PhotoCaptureResult { duplicated: false, block_count: count, preview, image_ref: rel }) |
| 101 | +} |
| 102 | + |
| 103 | +/// base64 解码 + 解码器验证(白名单由 image 解码器把关——capture_fragment 同模式)。 |
| 104 | +fn decode_image(image_b64: &str) -> Result<image::RgbImage, String> { |
| 105 | + let bytes = base64::engine::general_purpose::STANDARD |
| 106 | + .decode(image_b64.trim()) |
| 107 | + .map_err(|e| format!("图片 base64 解码失败: {}", e))?; |
| 108 | + if bytes.is_empty() { |
| 109 | + return Err("图片数据为空".to_string()); |
| 110 | + } |
| 111 | + if bytes.len() > MAX_PHOTO_BYTES { |
| 112 | + return Err(format!("图片超限(>{}MB)", MAX_PHOTO_BYTES / 1024 / 1024)); |
| 113 | + } |
| 114 | + let img = image::load_from_memory(&bytes) |
| 115 | + .map_err(|e| format!("图片解码失败(仅支持 PNG/JPEG 等常见格式): {}", e))?; |
| 116 | + let (w, h) = img.dimensions(); |
| 117 | + if w == 0 || h == 0 { |
| 118 | + return Err("图片尺寸非法".to_string()); |
| 119 | + } |
| 120 | + Ok(img.to_rgb8()) |
| 121 | +} |
| 122 | + |
| 123 | +/// 全图 OCR:缩至 960px 后识别,返回 (块, bbox 缩放比)。 |
| 124 | +/// @ai-context: 缩放比用于 bbox 反算回原图坐标系(TD-046 同思路——OCR 输入 |
| 125 | +/// 缩小后 bbox 处于缩小坐标系)。 |
| 126 | +fn ocr_blocks(img: &image::RgbImage, engines: &EnginePool) -> (Vec<crate::types::OcrBlock>, f32) { |
| 127 | + let (w, h) = img.dimensions(); |
| 128 | + let scale = if w > OCR_MAX_WIDTH { w as f32 / OCR_MAX_WIDTH as f32 } else { 1.0 }; |
| 129 | + let input = if w > OCR_MAX_WIDTH { |
| 130 | + let new_h = ((h as u64 * OCR_MAX_WIDTH as u64 / w as u64)).max(1) as u32; |
| 131 | + image::imageops::resize(img, OCR_MAX_WIDTH, new_h, image::imageops::FilterType::Triangle) |
| 132 | + } else { |
| 133 | + img.clone() |
| 134 | + }; |
| 135 | + match engines.recognize_image_timeout(input, crate::engine::OCR_REQUEST_TIMEOUT) { |
| 136 | + Ok(blocks) => (blocks, scale), |
| 137 | + Err(e) => { |
| 138 | + // OCR 失败/超时:图已保存,块数 0 诚实返回(下张截图重试语义) |
| 139 | + eprintln!("[Photo] OCR 失败(截图已保存,块数 0): {}", e); |
| 140 | + (Vec::new(), scale) |
| 141 | + } |
| 142 | + } |
| 143 | +} |
| 144 | + |
| 145 | +/// 单测独立文件(保持本文件 ≤300 行,AGENTS.md §3)。 |
| 146 | +#[cfg(test)] |
| 147 | +#[path = "photo_capture_tests.rs"] |
| 148 | +mod tests; |
0 commit comments