smriti/services/
document_detector.rs1#![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 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 if (aspect - 16.0 / 9.0).abs() < 0.20 && edge_density > 0.16 {
65 return ContentCategory::Presentation;
66 }
67
68 if edge_density > 0.15 && contrast > 0.20 && (0.45..=2.2).contains(&aspect) {
70 return ContentCategory::BusinessCard;
71 }
72
73 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}