oar_ocr_core/processors/decode.rs
1//! Text decoding utilities for OCR (Optical Character Recognition) systems.
2//!
3//! This module provides implementations for decoding text recognition results,
4//! particularly focused on CTC (Connectionist Temporal Classification) decoding.
5//! It includes structures and methods for converting model predictions into
6//! readable text strings with confidence scores.
7
8use rayon::prelude::*;
9use regex::Regex;
10use std::collections::HashMap;
11use std::sync::LazyLock;
12
13/// Decoded batch outputs along with positional metadata.
14pub type PositionedDecodeResult = (
15 Vec<String>,
16 Vec<f32>,
17 Vec<Vec<f32>>,
18 Vec<Vec<usize>>,
19 Vec<usize>,
20);
21
22static ALPHANUMERIC_REGEX: LazyLock<Regex> = LazyLock::new(|| {
23 Regex::new(r"[a-zA-Z0-9 :*./%+-]").expect("static regex: alphanumeric decoder pattern")
24});
25
26/// Argmax over a 1-D prediction row, returning `(index, value)`.
27///
28/// Contiguous rows (the common row-major case for the per-timestep logits) are
29/// routed through the SIMD kernel in [`crate::processors::simd`]; a scalar scan
30/// handles non-contiguous views. Tie-breaking matches [`Iterator::max_by`]
31/// (the last maximal index wins), so decoded output is unchanged.
32#[inline]
33fn argmax_row(row: ndarray::ArrayView1<f32>) -> Option<(usize, f32)> {
34 match row.as_slice() {
35 Some(slice) => crate::processors::simd::argmax(slice),
36 None => row
37 .iter()
38 .copied()
39 .enumerate()
40 .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)),
41 }
42}
43
44/// A base decoder for text recognition that handles character mapping and basic decoding operations.
45///
46/// This struct is responsible for converting model predictions into readable text strings.
47/// It maintains a character dictionary for mapping indices to characters and provides
48/// methods for decoding text with optional duplicate removal and confidence scoring.
49///
50/// # Fields
51/// * `reverse` - Flag indicating whether to reverse the text output
52/// * `dict` - A mapping from characters to their indices in the character list
53/// * `character` - A list of characters in the vocabulary, indexed by their position
54pub struct BaseRecLabelDecode {
55 reverse: bool,
56 dict: HashMap<char, usize>,
57 character: Vec<char>,
58}
59
60impl BaseRecLabelDecode {
61 /// Creates a new `BaseRecLabelDecode` instance.
62 ///
63 /// # Arguments
64 /// * `character_str` - An optional string containing the character vocabulary.
65 /// If None, a default alphanumeric character set is used.
66 /// * `use_space_char` - Whether to include a space character in the vocabulary.
67 ///
68 /// # Returns
69 /// A new `BaseRecLabelDecode` instance.
70 pub fn new(character_str: Option<&str>, use_space_char: bool) -> Self {
71 let mut character_list: Vec<char> = if let Some(chars) = character_str {
72 chars.chars().collect()
73 } else {
74 "0123456789abcdefghijklmnopqrstuvwxyz".chars().collect()
75 };
76
77 if use_space_char {
78 character_list.push(' ');
79 }
80
81 character_list = Self::add_special_char(character_list);
82
83 let mut dict = HashMap::new();
84 for (i, &char) in character_list.iter().enumerate() {
85 dict.insert(char, i);
86 }
87
88 Self {
89 reverse: false,
90 dict,
91 character: character_list,
92 }
93 }
94
95 /// Creates a new `BaseRecLabelDecode` instance from a list of strings.
96 ///
97 /// # Arguments
98 /// * `character_list` - An optional slice of strings containing the character vocabulary.
99 /// Only the first character of each string is used. If None, a default alphanumeric
100 /// character set is used.
101 /// * `use_space_char` - Whether to include a space character in the vocabulary.
102 ///
103 /// # Returns
104 /// A new `BaseRecLabelDecode` instance.
105 pub fn from_string_list(character_list: Option<&[String]>, use_space_char: bool) -> Self {
106 let mut chars: Vec<char> = if let Some(list) = character_list {
107 list.iter().filter_map(|s| s.chars().next()).collect()
108 } else {
109 "0123456789abcdefghijklmnopqrstuvwxyz".chars().collect()
110 };
111
112 if use_space_char {
113 chars.push(' ');
114 }
115
116 chars = Self::add_special_char(chars);
117
118 let mut dict = HashMap::new();
119 for (i, &char) in chars.iter().enumerate() {
120 dict.insert(char, i);
121 }
122
123 Self {
124 reverse: false,
125 dict,
126 character: chars,
127 }
128 }
129
130 /// Reverses the alphanumeric parts of a string while keeping non-alphanumeric parts in place.
131 ///
132 /// # Arguments
133 /// * `pred` - The input string to process.
134 ///
135 /// # Returns
136 /// A new string with alphanumeric parts reversed.
137 fn pred_reverse(&self, pred: &str) -> String {
138 let mut pred_re = Vec::new();
139 let mut c_current = String::new();
140
141 for c in pred.chars() {
142 if !ALPHANUMERIC_REGEX.is_match(&c.to_string()) {
143 if !c_current.is_empty() {
144 pred_re.push(c_current.clone());
145 c_current.clear();
146 }
147 pred_re.push(c.to_string());
148 } else {
149 c_current.push(c);
150 }
151 }
152
153 if !c_current.is_empty() {
154 pred_re.push(c_current);
155 }
156
157 pred_re.reverse();
158 pred_re.join("")
159 }
160
161 /// Adds special characters to the character list.
162 ///
163 /// This is a placeholder method that currently just returns the input list unchanged.
164 /// It can be overridden in subclasses to add special characters.
165 ///
166 /// # Arguments
167 /// * `character_list` - The input character list.
168 ///
169 /// # Returns
170 /// The character list with any special characters added.
171 fn add_special_char(character_list: Vec<char>) -> Vec<char> {
172 character_list
173 }
174
175 /// Gets a list of token indices that should be ignored during decoding.
176 ///
177 /// # Returns
178 /// A vector containing the indices of tokens to ignore.
179 fn get_ignored_tokens(&self) -> Vec<usize> {
180 vec![self.get_blank_idx()]
181 }
182
183 /// Decodes model predictions into text strings with confidence scores.
184 ///
185 /// # Arguments
186 /// * `text_index` - A slice of vectors containing the predicted character indices.
187 /// * `text_prob` - An optional slice of vectors containing the prediction probabilities.
188 /// * `is_remove_duplicate` - Whether to remove consecutive duplicate characters.
189 ///
190 /// # Returns
191 /// A vector of tuples, each containing a decoded text string and its confidence score.
192 pub fn decode(
193 &self,
194 text_index: &[Vec<usize>],
195 text_prob: Option<&[Vec<f32>]>,
196 is_remove_duplicate: bool,
197 ) -> Vec<(String, f32)> {
198 let mut result_list = Vec::new();
199 let ignored_tokens = self.get_ignored_tokens();
200
201 for (batch_idx, indices) in text_index.iter().enumerate() {
202 let mut selection = vec![true; indices.len()];
203
204 if is_remove_duplicate && indices.len() > 1 {
205 for i in 1..indices.len() {
206 if indices[i] == indices[i - 1] {
207 selection[i] = false;
208 }
209 }
210 }
211
212 for &ignored_token in &ignored_tokens {
213 for (i, &idx) in indices.iter().enumerate() {
214 if idx == ignored_token {
215 selection[i] = false;
216 }
217 }
218 }
219
220 let char_list: Vec<char> = indices
221 .iter()
222 .enumerate()
223 .filter(|(i, _)| selection[*i])
224 .filter_map(|(_, &text_id)| self.character.get(text_id).copied())
225 .collect();
226
227 let conf_list: Vec<f32> = if let Some(probs) = text_prob {
228 if batch_idx < probs.len() {
229 probs[batch_idx]
230 .iter()
231 .enumerate()
232 .filter(|(i, _)| *i < selection.len() && selection[*i])
233 .map(|(_, &prob)| prob)
234 .collect()
235 } else {
236 vec![1.0; char_list.len()]
237 }
238 } else {
239 vec![1.0; char_list.len()]
240 };
241
242 let conf_list = if conf_list.is_empty() {
243 vec![0.0]
244 } else {
245 conf_list
246 };
247
248 let mut text: String = char_list.iter().collect();
249
250 if self.reverse {
251 text = self.pred_reverse(&text);
252 }
253
254 let mean_conf = conf_list.iter().sum::<f32>() / conf_list.len() as f32;
255 result_list.push((text, mean_conf));
256 }
257
258 result_list
259 }
260
261 /// Applies the decoder to a tensor of model predictions.
262 ///
263 /// # Arguments
264 /// * `pred` - A 3D tensor containing the model predictions. Accepts any
265 /// `ndarray` storage (owned `Array3<f32>` or a zero-copy `ArrayView3<f32>`).
266 ///
267 /// # Returns
268 /// A tuple containing:
269 /// * A vector of decoded text strings
270 /// * A vector of confidence scores for each text string
271 pub fn apply(&self, pred: &ndarray::Array3<f32>) -> (Vec<String>, Vec<f32>) {
272 if pred.is_empty() {
273 return (Vec::new(), Vec::new());
274 }
275
276 let batch_size = pred.shape()[0];
277 let mut all_texts = Vec::new();
278 let mut all_scores = Vec::new();
279
280 for batch_idx in 0..batch_size {
281 let preds = pred.index_axis(ndarray::Axis(0), batch_idx);
282
283 let mut sequence_idx = Vec::new();
284 let mut sequence_prob = Vec::new();
285
286 for row in preds.outer_iter() {
287 if let Some((idx, prob)) = argmax_row(row) {
288 sequence_idx.push(idx);
289 sequence_prob.push(prob);
290 } else {
291 sequence_idx.push(0);
292 sequence_prob.push(0.0);
293 }
294 }
295
296 let text = self.decode(&[sequence_idx], Some(&[sequence_prob]), true);
297
298 for (t, score) in text {
299 all_texts.push(t);
300 all_scores.push(score);
301 }
302 }
303
304 (all_texts, all_scores)
305 }
306
307 /// Gets the index of the blank token.
308 ///
309 /// # Returns
310 /// The index of the blank token (always 0 in this base implementation).
311 fn get_blank_idx(&self) -> usize {
312 0
313 }
314}
315
316/// A decoder for CTC (Connectionist Temporal Classification) based text recognition models.
317///
318/// This struct extends `BaseRecLabelDecode` to provide specialized decoding for CTC models,
319/// which include a blank token that needs to be handled specially during decoding.
320///
321/// # Fields
322/// * `base` - The base decoder that handles character mapping and basic decoding operations
323/// * `blank_index` - The index of the blank token in the character vocabulary
324pub struct CTCLabelDecode {
325 base: BaseRecLabelDecode,
326 blank_index: usize,
327}
328
329impl std::fmt::Debug for CTCLabelDecode {
330 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
331 f.debug_struct("CTCLabelDecode")
332 .field("character_count", &self.base.character.len())
333 .field("reverse", &self.base.reverse)
334 .finish()
335 }
336}
337
338impl CTCLabelDecode {
339 /// Creates a new `CTCLabelDecode` instance.
340 ///
341 /// # Arguments
342 /// * `character_list` - An optional string containing the character vocabulary.
343 /// If None, a default alphanumeric character set is used.
344 /// * `use_space_char` - Whether to include a space character in the vocabulary.
345 ///
346 /// # Returns
347 /// A new `CTCLabelDecode` instance.
348 pub fn new(character_list: Option<&str>, use_space_char: bool) -> Self {
349 let mut base = BaseRecLabelDecode::new(character_list, use_space_char);
350
351 // Use null char for blank to distinguish from actual space
352 let mut new_character = vec!['\0'];
353 new_character.extend(base.character);
354
355 let mut new_dict = HashMap::new();
356 for (i, &char) in new_character.iter().enumerate() {
357 new_dict.insert(char, i);
358 }
359
360 base.character = new_character;
361 base.dict = new_dict;
362
363 let blank_index = 0;
364
365 Self { base, blank_index }
366 }
367
368 /// Creates a new `CTCLabelDecode` instance from a list of strings.
369 ///
370 /// # Arguments
371 /// * `character_list` - An optional slice of strings containing the character vocabulary.
372 /// Only the first character of each string is used. If None, a default alphanumeric
373 /// character set is used.
374 /// * `use_space_char` - Whether to include a space character in the vocabulary.
375 /// * `has_explicit_blank` - Whether the character list already includes a blank token.
376 ///
377 /// # Returns
378 /// A new `CTCLabelDecode` instance.
379 pub fn from_string_list(
380 character_list: Option<&[String]>,
381 use_space_char: bool,
382 has_explicit_blank: bool,
383 ) -> Self {
384 if has_explicit_blank {
385 let base = BaseRecLabelDecode::from_string_list(character_list, use_space_char);
386 Self {
387 base,
388 blank_index: 0,
389 }
390 } else {
391 let mut base = BaseRecLabelDecode::from_string_list(character_list, use_space_char);
392
393 // Use null char for blank to distinguish from actual space
394 let mut new_character = vec!['\0'];
395 new_character.extend(base.character);
396
397 let mut new_dict = HashMap::new();
398 for (i, &char) in new_character.iter().enumerate() {
399 new_dict.insert(char, i);
400 }
401
402 base.character = new_character;
403 base.dict = new_dict;
404
405 Self {
406 base,
407 blank_index: 0,
408 }
409 }
410 }
411
412 /// Gets the index of the blank token.
413 ///
414 /// # Returns
415 /// The index of the blank token.
416 pub fn get_blank_index(&self) -> usize {
417 self.blank_index
418 }
419
420 /// Gets the character list used by this decoder.
421 ///
422 /// # Returns
423 /// A slice containing the characters in the vocabulary.
424 pub fn get_character_list(&self) -> &[char] {
425 &self.base.character
426 }
427
428 /// Gets the number of characters in the vocabulary.
429 ///
430 /// # Returns
431 /// The number of characters in the vocabulary.
432 pub fn get_character_count(&self) -> usize {
433 self.base.character.len()
434 }
435
436 /// Applies the CTC decoder to a tensor of model predictions with character position tracking.
437 ///
438 /// This method handles the special requirements of CTC decoding and additionally tracks
439 /// the timestep positions of each character for word box generation.
440 ///
441 /// # Arguments
442 /// * `pred` - A 3D tensor containing the model predictions. Accepts any
443 /// `ndarray` storage (owned `Array3<f32>` or a zero-copy `ArrayView3<f32>`).
444 ///
445 /// # Returns
446 /// A tuple containing:
447 /// * A vector of decoded text strings
448 /// * A vector of confidence scores for each text string
449 /// * A vector of character positions (normalized 0.0-1.0) for each text string
450 /// * A vector of column indices for each character in each text string
451 /// * A vector of sequence lengths (total columns) for each text string
452 pub fn apply_with_positions<S>(
453 &self,
454 pred: &ndarray::ArrayBase<S, ndarray::Ix3>,
455 ) -> PositionedDecodeResult
456 where
457 S: ndarray::Data<Elem = f32> + Sync,
458 {
459 if pred.is_empty() {
460 return (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new());
461 }
462
463 let batch_size = pred.shape()[0];
464
465 type PerItem = (String, f32, Vec<f32>, Vec<usize>, usize);
466 let per: Vec<PerItem> = (0..batch_size)
467 .into_par_iter()
468 .map(|batch_idx| {
469 let preds = pred.index_axis(ndarray::Axis(0), batch_idx);
470 let seq_len_usize = preds.shape()[0];
471 let seq_len = seq_len_usize as f32;
472
473 let mut sequence_idx = Vec::with_capacity(seq_len_usize);
474 let mut sequence_prob = Vec::with_capacity(seq_len_usize);
475
476 for row in preds.outer_iter() {
477 if let Some((idx, prob)) = argmax_row(row) {
478 sequence_idx.push(idx);
479 sequence_prob.push(prob);
480 } else {
481 sequence_idx.push(self.blank_index);
482 sequence_prob.push(0.0);
483 }
484 }
485
486 // Single CTC collapse pass (timestep == sequence position): drop
487 // blanks and consecutive duplicates and map to glyphs in one go,
488 // avoiding the `selection` scratch vector and two extra passes.
489 // `prev_idx` is updated on every step (blanks included), so dedup
490 // runs on the raw indices exactly as before. Pushing char/prob/
491 // timestep together keeps an out-of-vocab index from desyncing
492 // `char_list` from `char_positions` and corrupting word boxes.
493 let mut filtered_prob = Vec::with_capacity(sequence_idx.len());
494 let mut filtered_timesteps = Vec::with_capacity(sequence_idx.len());
495 let mut char_list: Vec<char> = Vec::with_capacity(sequence_idx.len());
496 let mut prev_idx = self.blank_index;
497 for (i, &idx) in sequence_idx.iter().enumerate() {
498 if idx != self.blank_index
499 && idx != prev_idx
500 && let Some(&ch) = self.base.character.get(idx)
501 {
502 char_list.push(ch);
503 filtered_prob.push(sequence_prob[i]);
504 filtered_timesteps.push(i);
505 }
506 prev_idx = idx;
507 }
508
509 let mean_conf = if filtered_prob.is_empty() {
510 0.0
511 } else {
512 filtered_prob.iter().sum::<f32>() / filtered_prob.len() as f32
513 };
514
515 // Calculate normalized character positions (0.0 to 1.0)
516 let char_positions: Vec<f32> = filtered_timesteps
517 .iter()
518 .map(|×tep| timestep as f32 / seq_len)
519 .collect();
520
521 let text: String = char_list.iter().collect();
522 (
523 text,
524 mean_conf,
525 char_positions,
526 filtered_timesteps,
527 seq_len_usize,
528 )
529 })
530 .collect();
531
532 let mut all_texts = Vec::with_capacity(batch_size);
533 let mut all_scores = Vec::with_capacity(batch_size);
534 let mut all_positions = Vec::with_capacity(batch_size);
535 let mut all_col_indices = Vec::with_capacity(batch_size);
536 let mut all_seq_lengths = Vec::with_capacity(batch_size);
537 for (text, score, pos, cols, seq_len) in per {
538 all_texts.push(text);
539 all_scores.push(score);
540 all_positions.push(pos);
541 all_col_indices.push(cols);
542 all_seq_lengths.push(seq_len);
543 }
544
545 (
546 all_texts,
547 all_scores,
548 all_positions,
549 all_col_indices,
550 all_seq_lengths,
551 )
552 }
553
554 /// Applies the CTC decoder to a tensor of model predictions.
555 ///
556 /// This method handles the special requirements of CTC decoding:
557 /// 1. Removing blank tokens
558 /// 2. Removing consecutive duplicate characters
559 /// 3. Converting indices to characters
560 /// 4. Calculating confidence scores
561 ///
562 /// # Arguments
563 /// * `pred` - A 3D tensor containing the model predictions. Accepts any
564 /// `ndarray` storage (owned `Array3<f32>` or a zero-copy `ArrayView3<f32>`).
565 ///
566 /// # Returns
567 /// A tuple containing:
568 /// * A vector of decoded text strings
569 /// * A vector of confidence scores for each text string
570 pub fn apply<S>(&self, pred: &ndarray::ArrayBase<S, ndarray::Ix3>) -> (Vec<String>, Vec<f32>)
571 where
572 S: ndarray::Data<Elem = f32> + Sync,
573 {
574 if pred.is_empty() {
575 return (Vec::new(), Vec::new());
576 }
577
578 let batch_size = pred.shape()[0];
579
580 // Decode each sequence in the batch independently. The argmax over the
581 // (large) vocab dimension for every timestep dominates this work, so we
582 // fan the per-sequence loop out across rayon — order is preserved by
583 // collecting into an indexed Vec.
584 let (all_texts, all_scores): (Vec<String>, Vec<f32>) = (0..batch_size)
585 .into_par_iter()
586 .map(|batch_idx| {
587 let preds = pred.index_axis(ndarray::Axis(0), batch_idx);
588 let seq_len_usize = preds.shape()[0];
589
590 let mut sequence_idx = Vec::with_capacity(seq_len_usize);
591 let mut sequence_prob = Vec::with_capacity(seq_len_usize);
592
593 for row in preds.outer_iter() {
594 if let Some((idx, prob)) = argmax_row(row) {
595 sequence_idx.push(idx);
596 sequence_prob.push(prob);
597 } else {
598 sequence_idx.push(self.blank_index);
599 sequence_prob.push(0.0);
600 }
601 }
602
603 // Single CTC collapse pass: drop blanks and consecutive duplicates
604 // and map to glyphs in one go, avoiding the `selection` scratch
605 // vector and two extra passes. `prev_idx` is updated on every step
606 // (blanks included), so dedup runs on the raw indices exactly as
607 // before. Only count a prob when its glyph lands in `text`, else an
608 // out-of-vocab index would inflate `mean_conf`.
609 let mut filtered_prob = Vec::with_capacity(sequence_idx.len());
610 let mut text = String::with_capacity(sequence_idx.len());
611 let mut prev_idx = self.blank_index;
612 for (i, &idx) in sequence_idx.iter().enumerate() {
613 if idx != self.blank_index
614 && idx != prev_idx
615 && let Some(&ch) = self.base.character.get(idx)
616 {
617 text.push(ch);
618 filtered_prob.push(sequence_prob[i]);
619 }
620 prev_idx = idx;
621 }
622
623 let mean_conf = if filtered_prob.is_empty() {
624 0.0
625 } else {
626 filtered_prob.iter().sum::<f32>() / filtered_prob.len() as f32
627 };
628
629 (text, mean_conf)
630 })
631 .unzip();
632
633 (all_texts, all_scores)
634 }
635}