Skip to main content

smriti/services/
document_detector.rs

1//! Heuristic document detector (Stage 1 fast filter).
2//!
3//! Temporarily disabled - call sites in ocr_processor.rs are commented out.
4//! Kept compilable so we can revisit with a learned classifier later.
5#![allow(dead_code)]
6
7use std::path::Path;
8
9use image::{DynamicImage, GenericImageView};
10
11use crate::models::ContentCategory;
12
13pub struct DocumentDetector;
14
15impl DocumentDetector {
16    pub fn classify(image: &DynamicImage, file_path: &str) -> ContentCategory {
17        let file_name = Path::new(file_path)
18            .file_name()
19            .and_then(|n| n.to_str())
20            .unwrap_or_default()
21            .to_lowercase();
22
23        if file_name.contains("screenshot")
24            || file_name.starts_with("img_") && file_name.contains("wa")
25        {
26            return ContentCategory::Screenshot;
27        }
28        if file_name.contains("receipt")
29            || file_name.contains("invoice")
30            || file_name.contains("bill")
31        {
32            return ContentCategory::Receipt;
33        }
34        if file_name.contains("business card") || file_name.contains("bizcard") {
35            return ContentCategory::BusinessCard;
36        }
37        if file_name.contains("slide") || file_name.contains("ppt") || file_name.contains("deck") {
38            return ContentCategory::Presentation;
39        }
40        if file_name.contains("whiteboard") {
41            return ContentCategory::Whiteboard;
42        }
43
44        let (w, h) = image.dimensions();
45        if w == 0 || h == 0 {
46            return ContentCategory::Photo;
47        }
48
49        let aspect = w as f32 / h as f32;
50        let resized = image
51            .resize(512, 512, image::imageops::FilterType::Triangle)
52            .to_luma8();
53
54        let edge_density = Self::edge_density(&resized);
55        let contrast = Self::contrast_spread(&resized);
56
57        // Screenshot-like: precise edges + wide/phone aspect.
58        let looks_like_screen = edge_density > 0.20 && !(0.72..=1.45).contains(&aspect);
59        if looks_like_screen {
60            return ContentCategory::Screenshot;
61        }
62
63        // Presentation-like: 16:9-ish + high edge structure.
64        if (aspect - 16.0 / 9.0).abs() < 0.20 && edge_density > 0.16 {
65            return ContentCategory::Presentation;
66        }
67
68        // Card-like: compact aspect and high-contrast text blocks.
69        if edge_density > 0.15 && contrast > 0.20 && (0.45..=2.2).contains(&aspect) {
70            return ContentCategory::BusinessCard;
71        }
72
73        // Document-like: lots of text edges and moderate/high contrast.
74        if edge_density > 0.14 && contrast > 0.18 {
75            return ContentCategory::Document;
76        }
77
78        ContentCategory::Photo
79    }
80
81    pub fn classify_with_text_hints(
82        image: &DynamicImage,
83        file_path: &str,
84        ocr_text: Option<&str>,
85    ) -> ContentCategory {
86        let mut category = Self::classify(image, file_path);
87
88        let text = ocr_text.unwrap_or_default().to_lowercase();
89        if !text.is_empty() {
90            let has_phone =
91                text.contains("phone") || text.contains("mobile") || text.contains("tel");
92            let has_email = text.contains('@') || text.contains("email");
93            if has_phone && has_email {
94                return ContentCategory::BusinessCard;
95            }
96
97            let bullet_markers = text.matches("\n-").count()
98                + text.matches("\n*").count()
99                + text.matches("\n•").count();
100            if bullet_markers >= 4 {
101                category = ContentCategory::Presentation;
102            }
103        }
104
105        category
106    }
107
108    fn edge_density(gray: &image::GrayImage) -> f32 {
109        let (w, h) = gray.dimensions();
110        if w < 3 || h < 3 {
111            return 0.0;
112        }
113
114        let mut edge_count = 0u64;
115        let mut total = 0u64;
116
117        for y in 1..(h - 1) {
118            for x in 1..(w - 1) {
119                let l = gray.get_pixel(x - 1, y)[0] as i32;
120                let r = gray.get_pixel(x + 1, y)[0] as i32;
121                let u = gray.get_pixel(x, y - 1)[0] as i32;
122                let d = gray.get_pixel(x, y + 1)[0] as i32;
123                let grad = (r - l).abs() + (d - u).abs();
124                if grad > 48 {
125                    edge_count += 1;
126                }
127                total += 1;
128            }
129        }
130
131        if total == 0 {
132            0.0
133        } else {
134            edge_count as f32 / total as f32
135        }
136    }
137
138    fn contrast_spread(gray: &image::GrayImage) -> f32 {
139        let mut min_v = 255u8;
140        let mut max_v = 0u8;
141        for p in gray.pixels() {
142            let v = p[0];
143            if v < min_v {
144                min_v = v;
145            }
146            if v > max_v {
147                max_v = v;
148            }
149        }
150        (max_v.saturating_sub(min_v) as f32) / 255.0
151    }
152}