Skip to main content

smriti/services/
ocr_processor.rs

1//! OCR processor scaffolding (Stage 2 deep analysis architecture).
2//!
3//! Stage 1 uses heuristic-only document detection. This service provides
4//! progress/cancellation architecture so OCR can be plugged in with ONNX models.
5
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::Arc;
8
9use crate::db::Database;
10use crate::db::DocumentRepo;
11
12#[derive(Debug, Clone)]
13pub struct OcrProgress {
14    pub processed: usize,
15    pub total: usize,
16    pub documents_found: usize,
17}
18
19pub struct OcrProcessor;
20
21impl OcrProcessor {
22    /// Mark every unprocessed photo as `Photo` and tick the
23    /// content-categorisation flag. The earlier heuristic classifier
24    /// (edge density + aspect ratio + filename keywords) produced too
25    /// many false positives and was retired; this scaffolding stays so
26    /// a learned classifier can be plugged in later without re-piping
27    /// the progress / cancellation channels.
28    pub fn process_stage1_heuristics(
29        drive_path: &std::path::Path,
30        progress_tx: Option<async_channel::Sender<OcrProgress>>,
31        cancel_flag: Option<Arc<AtomicBool>>,
32    ) -> Result<usize, String> {
33        let db = Database::open_for_drive(drive_path)
34            .map_err(|e| format!("Failed to open database: {}", e))?;
35        let repo = DocumentRepo::new(&db.conn);
36
37        let targets = repo
38            .get_unprocessed_for_document_analysis(100000)
39            .map_err(|e| format!("Failed to load analysis targets: {}", e))?;
40
41        let total = targets.len();
42        if total == 0 {
43            return Ok(0);
44        }
45
46        for (idx, (photo_id, _rel_path, _orientation)) in targets.iter().enumerate() {
47            if let Some(ref flag) = cancel_flag {
48                if flag.load(Ordering::Relaxed) {
49                    return Ok(0);
50                }
51            }
52
53            let _ = repo
54                .update_content_category(*photo_id, crate::models::ContentCategory::Photo.as_str());
55            let _ = repo.update_ocr_metadata(*photo_id, None, None, true);
56
57            if let Some(ref tx) = progress_tx {
58                let _ = tx.try_send(OcrProgress {
59                    processed: idx + 1,
60                    total,
61                    documents_found: 0,
62                });
63            }
64        }
65
66        Ok(0)
67    }
68}