Skip to main content

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
22/// Compact result of reducing CTC logits over the vocabulary dimension.
23///
24/// Unlike the original `(batch, time, vocab)` logits, this only retains one
25/// token index and confidence per timestep, so it is cheap to move beyond the
26/// lifetime of an ONNX Runtime output buffer.
27#[derive(Debug, PartialEq)]
28pub(crate) struct CTCArgmaxOutput {
29    batch_size: usize,
30    sequence_length: usize,
31    indices: Vec<usize>,
32    probabilities: Vec<f32>,
33}
34
35static ALPHANUMERIC_REGEX: LazyLock<Regex> = LazyLock::new(|| {
36    Regex::new(r"[a-zA-Z0-9 :*./%+-]").expect("static regex: alphanumeric decoder pattern")
37});
38
39/// Argmax over a 1-D prediction row, returning `(index, value)`.
40///
41/// Contiguous rows (the common row-major case for the per-timestep logits) are
42/// routed through the SIMD kernel in [`crate::processors::simd`]; a scalar scan
43/// handles non-contiguous views. Tie-breaking matches [`Iterator::max_by`]
44/// (the last maximal index wins), so decoded output is unchanged.
45#[inline]
46fn argmax_row(row: ndarray::ArrayView1<f32>) -> Option<(usize, f32)> {
47    match row.as_slice() {
48        Some(slice) => crate::processors::simd::argmax(slice),
49        None => row
50            .iter()
51            .copied()
52            .enumerate()
53            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)),
54    }
55}
56
57/// A base decoder for text recognition that handles character mapping and basic decoding operations.
58///
59/// This struct is responsible for converting model predictions into readable text strings.
60/// It maintains a character dictionary for mapping indices to characters and provides
61/// methods for decoding text with optional duplicate removal and confidence scoring.
62///
63/// # Fields
64/// * `reverse` - Flag indicating whether to reverse the text output
65/// * `dict` - A mapping from characters to their indices in the character list
66/// * `character` - A list of characters in the vocabulary, indexed by their position
67pub struct BaseRecLabelDecode {
68    reverse: bool,
69    dict: HashMap<char, usize>,
70    character: Vec<char>,
71}
72
73impl BaseRecLabelDecode {
74    /// Creates a new `BaseRecLabelDecode` instance.
75    ///
76    /// # Arguments
77    /// * `character_str` - An optional string containing the character vocabulary.
78    ///   If None, a default alphanumeric character set is used.
79    /// * `use_space_char` - Whether to include a space character in the vocabulary.
80    ///
81    /// # Returns
82    /// A new `BaseRecLabelDecode` instance.
83    pub fn new(character_str: Option<&str>, use_space_char: bool) -> Self {
84        let mut character_list: Vec<char> = if let Some(chars) = character_str {
85            chars.chars().collect()
86        } else {
87            "0123456789abcdefghijklmnopqrstuvwxyz".chars().collect()
88        };
89
90        if use_space_char {
91            character_list.push(' ');
92        }
93
94        character_list = Self::add_special_char(character_list);
95
96        let mut dict = HashMap::new();
97        for (i, &char) in character_list.iter().enumerate() {
98            dict.insert(char, i);
99        }
100
101        Self {
102            reverse: false,
103            dict,
104            character: character_list,
105        }
106    }
107
108    /// Creates a new `BaseRecLabelDecode` instance from a list of strings.
109    ///
110    /// # Arguments
111    /// * `character_list` - An optional slice of strings containing the character vocabulary.
112    ///   Only the first character of each string is used. If None, a default alphanumeric
113    ///   character set is used.
114    /// * `use_space_char` - Whether to include a space character in the vocabulary.
115    ///
116    /// # Returns
117    /// A new `BaseRecLabelDecode` instance.
118    pub fn from_string_list(character_list: Option<&[String]>, use_space_char: bool) -> Self {
119        let mut chars: Vec<char> = if let Some(list) = character_list {
120            list.iter().filter_map(|s| s.chars().next()).collect()
121        } else {
122            "0123456789abcdefghijklmnopqrstuvwxyz".chars().collect()
123        };
124
125        if use_space_char {
126            chars.push(' ');
127        }
128
129        chars = Self::add_special_char(chars);
130
131        let mut dict = HashMap::new();
132        for (i, &char) in chars.iter().enumerate() {
133            dict.insert(char, i);
134        }
135
136        Self {
137            reverse: false,
138            dict,
139            character: chars,
140        }
141    }
142
143    /// Reverses the alphanumeric parts of a string while keeping non-alphanumeric parts in place.
144    ///
145    /// # Arguments
146    /// * `pred` - The input string to process.
147    ///
148    /// # Returns
149    /// A new string with alphanumeric parts reversed.
150    fn pred_reverse(&self, pred: &str) -> String {
151        let mut pred_re = Vec::new();
152        let mut c_current = String::new();
153
154        for c in pred.chars() {
155            if !ALPHANUMERIC_REGEX.is_match(&c.to_string()) {
156                if !c_current.is_empty() {
157                    pred_re.push(c_current.clone());
158                    c_current.clear();
159                }
160                pred_re.push(c.to_string());
161            } else {
162                c_current.push(c);
163            }
164        }
165
166        if !c_current.is_empty() {
167            pred_re.push(c_current);
168        }
169
170        pred_re.reverse();
171        pred_re.join("")
172    }
173
174    /// Adds special characters to the character list.
175    ///
176    /// This is a placeholder method that currently just returns the input list unchanged.
177    /// It can be overridden in subclasses to add special characters.
178    ///
179    /// # Arguments
180    /// * `character_list` - The input character list.
181    ///
182    /// # Returns
183    /// The character list with any special characters added.
184    fn add_special_char(character_list: Vec<char>) -> Vec<char> {
185        character_list
186    }
187
188    /// Gets a list of token indices that should be ignored during decoding.
189    ///
190    /// # Returns
191    /// A vector containing the indices of tokens to ignore.
192    fn get_ignored_tokens(&self) -> Vec<usize> {
193        vec![self.get_blank_idx()]
194    }
195
196    /// Decodes model predictions into text strings with confidence scores.
197    ///
198    /// # Arguments
199    /// * `text_index` - A slice of vectors containing the predicted character indices.
200    /// * `text_prob` - An optional slice of vectors containing the prediction probabilities.
201    /// * `is_remove_duplicate` - Whether to remove consecutive duplicate characters.
202    ///
203    /// # Returns
204    /// A vector of tuples, each containing a decoded text string and its confidence score.
205    pub fn decode(
206        &self,
207        text_index: &[Vec<usize>],
208        text_prob: Option<&[Vec<f32>]>,
209        is_remove_duplicate: bool,
210    ) -> Vec<(String, f32)> {
211        let mut result_list = Vec::new();
212        let ignored_tokens = self.get_ignored_tokens();
213
214        for (batch_idx, indices) in text_index.iter().enumerate() {
215            let mut selection = vec![true; indices.len()];
216
217            if is_remove_duplicate && indices.len() > 1 {
218                for i in 1..indices.len() {
219                    if indices[i] == indices[i - 1] {
220                        selection[i] = false;
221                    }
222                }
223            }
224
225            for &ignored_token in &ignored_tokens {
226                for (i, &idx) in indices.iter().enumerate() {
227                    if idx == ignored_token {
228                        selection[i] = false;
229                    }
230                }
231            }
232
233            let char_list: Vec<char> = indices
234                .iter()
235                .enumerate()
236                .filter(|(i, _)| selection[*i])
237                .filter_map(|(_, &text_id)| self.character.get(text_id).copied())
238                .collect();
239
240            let conf_list: Vec<f32> = if let Some(probs) = text_prob {
241                if batch_idx < probs.len() {
242                    probs[batch_idx]
243                        .iter()
244                        .enumerate()
245                        .filter(|(i, _)| *i < selection.len() && selection[*i])
246                        .map(|(_, &prob)| prob)
247                        .collect()
248                } else {
249                    vec![1.0; char_list.len()]
250                }
251            } else {
252                vec![1.0; char_list.len()]
253            };
254
255            let conf_list = if conf_list.is_empty() {
256                vec![0.0]
257            } else {
258                conf_list
259            };
260
261            let mut text: String = char_list.iter().collect();
262
263            if self.reverse {
264                text = self.pred_reverse(&text);
265            }
266
267            let mean_conf = conf_list.iter().sum::<f32>() / conf_list.len() as f32;
268            result_list.push((text, mean_conf));
269        }
270
271        result_list
272    }
273
274    /// Applies the decoder to a tensor of model predictions.
275    ///
276    /// # Arguments
277    /// * `pred` - A 3D tensor containing the model predictions. Accepts any
278    ///   `ndarray` storage (owned `Array3<f32>` or a zero-copy `ArrayView3<f32>`).
279    ///
280    /// # Returns
281    /// A tuple containing:
282    /// * A vector of decoded text strings
283    /// * A vector of confidence scores for each text string
284    pub fn apply(&self, pred: &ndarray::Array3<f32>) -> (Vec<String>, Vec<f32>) {
285        if pred.is_empty() {
286            return (Vec::new(), Vec::new());
287        }
288
289        let batch_size = pred.shape()[0];
290        let mut all_texts = Vec::new();
291        let mut all_scores = Vec::new();
292
293        for batch_idx in 0..batch_size {
294            let preds = pred.index_axis(ndarray::Axis(0), batch_idx);
295
296            let mut sequence_idx = Vec::new();
297            let mut sequence_prob = Vec::new();
298
299            for row in preds.outer_iter() {
300                if let Some((idx, prob)) = argmax_row(row) {
301                    sequence_idx.push(idx);
302                    sequence_prob.push(prob);
303                } else {
304                    sequence_idx.push(0);
305                    sequence_prob.push(0.0);
306                }
307            }
308
309            let text = self.decode(&[sequence_idx], Some(&[sequence_prob]), true);
310
311            for (t, score) in text {
312                all_texts.push(t);
313                all_scores.push(score);
314            }
315        }
316
317        (all_texts, all_scores)
318    }
319
320    /// Gets the index of the blank token.
321    ///
322    /// # Returns
323    /// The index of the blank token (always 0 in this base implementation).
324    fn get_blank_idx(&self) -> usize {
325        0
326    }
327}
328
329/// A decoder for CTC (Connectionist Temporal Classification) based text recognition models.
330///
331/// This struct extends `BaseRecLabelDecode` to provide specialized decoding for CTC models,
332/// which include a blank token that needs to be handled specially during decoding.
333///
334/// # Fields
335/// * `base` - The base decoder that handles character mapping and basic decoding operations
336/// * `blank_index` - The index of the blank token in the character vocabulary
337pub struct CTCLabelDecode {
338    base: BaseRecLabelDecode,
339    blank_index: usize,
340}
341
342impl std::fmt::Debug for CTCLabelDecode {
343    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
344        f.debug_struct("CTCLabelDecode")
345            .field("character_count", &self.base.character.len())
346            .field("reverse", &self.base.reverse)
347            .finish()
348    }
349}
350
351impl CTCLabelDecode {
352    /// Creates a new `CTCLabelDecode` instance.
353    ///
354    /// # Arguments
355    /// * `character_list` - An optional string containing the character vocabulary.
356    ///   If None, a default alphanumeric character set is used.
357    /// * `use_space_char` - Whether to include a space character in the vocabulary.
358    ///
359    /// # Returns
360    /// A new `CTCLabelDecode` instance.
361    pub fn new(character_list: Option<&str>, use_space_char: bool) -> Self {
362        let mut base = BaseRecLabelDecode::new(character_list, use_space_char);
363
364        // Use null char for blank to distinguish from actual space
365        let mut new_character = vec!['\0'];
366        new_character.extend(base.character);
367
368        let mut new_dict = HashMap::new();
369        for (i, &char) in new_character.iter().enumerate() {
370            new_dict.insert(char, i);
371        }
372
373        base.character = new_character;
374        base.dict = new_dict;
375
376        let blank_index = 0;
377
378        Self { base, blank_index }
379    }
380
381    /// Creates a new `CTCLabelDecode` instance from a list of strings.
382    ///
383    /// # Arguments
384    /// * `character_list` - An optional slice of strings containing the character vocabulary.
385    ///   Only the first character of each string is used. If None, a default alphanumeric
386    ///   character set is used.
387    /// * `use_space_char` - Whether to include a space character in the vocabulary.
388    /// * `has_explicit_blank` - Whether the character list already includes a blank token.
389    ///
390    /// # Returns
391    /// A new `CTCLabelDecode` instance.
392    pub fn from_string_list(
393        character_list: Option<&[String]>,
394        use_space_char: bool,
395        has_explicit_blank: bool,
396    ) -> Self {
397        if has_explicit_blank {
398            let base = BaseRecLabelDecode::from_string_list(character_list, use_space_char);
399            Self {
400                base,
401                blank_index: 0,
402            }
403        } else {
404            let mut base = BaseRecLabelDecode::from_string_list(character_list, use_space_char);
405
406            // Use null char for blank to distinguish from actual space
407            let mut new_character = vec!['\0'];
408            new_character.extend(base.character);
409
410            let mut new_dict = HashMap::new();
411            for (i, &char) in new_character.iter().enumerate() {
412                new_dict.insert(char, i);
413            }
414
415            base.character = new_character;
416            base.dict = new_dict;
417
418            Self {
419                base,
420                blank_index: 0,
421            }
422        }
423    }
424
425    /// Gets the index of the blank token.
426    ///
427    /// # Returns
428    /// The index of the blank token.
429    pub fn get_blank_index(&self) -> usize {
430        self.blank_index
431    }
432
433    /// Gets the character list used by this decoder.
434    ///
435    /// # Returns
436    /// A slice containing the characters in the vocabulary.
437    pub fn get_character_list(&self) -> &[char] {
438        &self.base.character
439    }
440
441    /// Gets the number of characters in the vocabulary.
442    ///
443    /// # Returns
444    /// The number of characters in the vocabulary.
445    pub fn get_character_count(&self) -> usize {
446        self.base.character.len()
447    }
448
449    /// Reduces `(batch, time, vocab)` logits to one token and confidence per
450    /// timestep. This is the only part of CTC decoding that must inspect the
451    /// large logits buffer.
452    pub(crate) fn argmax_predictions<S>(
453        &self,
454        pred: &ndarray::ArrayBase<S, ndarray::Ix3>,
455    ) -> CTCArgmaxOutput
456    where
457        S: ndarray::Data<Elem = f32> + Sync,
458    {
459        let [batch_size, sequence_length, vocab_size] = pred
460            .shape()
461            .try_into()
462            .expect("CTC predictions always have three dimensions");
463        let row_count = batch_size * sequence_length;
464
465        // Preserve the public decoder's historical empty-tensor behavior: no
466        // batch entries are returned when any dimension is zero.
467        if pred.is_empty() {
468            return CTCArgmaxOutput {
469                batch_size: 0,
470                sequence_length: 0,
471                indices: Vec::new(),
472                probabilities: Vec::new(),
473            };
474        }
475
476        // ORT outputs are contiguous, so the hot path can fan all timesteps out
477        // across rayon, including batch-size 1. Keep a non-contiguous fallback
478        // for callers of the public ndarray-based decoder methods.
479        let (indices, probabilities): (Vec<usize>, Vec<f32>) = if let Some(data) = pred.as_slice() {
480            data.par_chunks_exact(vocab_size)
481                .map(|row| crate::processors::simd::argmax(row).unwrap_or((self.blank_index, 0.0)))
482                .unzip()
483        } else {
484            (0..row_count)
485                .into_par_iter()
486                .map(|row_idx| {
487                    let batch_idx = row_idx / sequence_length;
488                    let time_idx = row_idx % sequence_length;
489                    argmax_row(pred.slice(ndarray::s![batch_idx, time_idx, ..]))
490                        .unwrap_or((self.blank_index, 0.0))
491                })
492                .unzip()
493        };
494
495        CTCArgmaxOutput {
496            batch_size,
497            sequence_length,
498            indices,
499            probabilities,
500        }
501    }
502
503    /// Performs CTC collapse and text construction from compact argmax data.
504    /// This no longer needs access to the logits or the inference session.
505    pub(crate) fn decode_argmax(&self, argmax: &CTCArgmaxOutput) -> (Vec<String>, Vec<f32>) {
506        let (all_texts, all_scores): (Vec<String>, Vec<f32>) = (0..argmax.batch_size)
507            .into_par_iter()
508            .map(|batch_idx| {
509                let start = batch_idx * argmax.sequence_length;
510                let end = start + argmax.sequence_length;
511                let sequence_idx = &argmax.indices[start..end];
512                let sequence_prob = &argmax.probabilities[start..end];
513
514                let mut filtered_prob = Vec::with_capacity(argmax.sequence_length);
515                let mut text = String::with_capacity(argmax.sequence_length);
516                let mut prev_idx = self.blank_index;
517                for (i, &idx) in sequence_idx.iter().enumerate() {
518                    if idx != self.blank_index
519                        && idx != prev_idx
520                        && let Some(&ch) = self.base.character.get(idx)
521                    {
522                        text.push(ch);
523                        filtered_prob.push(sequence_prob[i]);
524                    }
525                    prev_idx = idx;
526                }
527
528                let mean_conf = if filtered_prob.is_empty() {
529                    0.0
530                } else {
531                    filtered_prob.iter().sum::<f32>() / filtered_prob.len() as f32
532                };
533
534                (text, mean_conf)
535            })
536            .unzip();
537
538        (all_texts, all_scores)
539    }
540
541    /// Performs CTC collapse while retaining character timestep positions.
542    /// This no longer needs access to the logits or the inference session.
543    pub(crate) fn decode_argmax_with_positions(
544        &self,
545        argmax: &CTCArgmaxOutput,
546    ) -> PositionedDecodeResult {
547        type PerItem = (String, f32, Vec<f32>, Vec<usize>, usize);
548        let per: Vec<PerItem> = (0..argmax.batch_size)
549            .into_par_iter()
550            .map(|batch_idx| {
551                let start = batch_idx * argmax.sequence_length;
552                let end = start + argmax.sequence_length;
553                let sequence_idx = &argmax.indices[start..end];
554                let sequence_prob = &argmax.probabilities[start..end];
555
556                let mut filtered_prob = Vec::with_capacity(argmax.sequence_length);
557                let mut filtered_timesteps = Vec::with_capacity(argmax.sequence_length);
558                let mut char_list = Vec::with_capacity(argmax.sequence_length);
559                let mut prev_idx = self.blank_index;
560                for (i, &idx) in sequence_idx.iter().enumerate() {
561                    if idx != self.blank_index
562                        && idx != prev_idx
563                        && let Some(&ch) = self.base.character.get(idx)
564                    {
565                        char_list.push(ch);
566                        filtered_prob.push(sequence_prob[i]);
567                        filtered_timesteps.push(i);
568                    }
569                    prev_idx = idx;
570                }
571
572                let mean_conf = if filtered_prob.is_empty() {
573                    0.0
574                } else {
575                    filtered_prob.iter().sum::<f32>() / filtered_prob.len() as f32
576                };
577                let seq_len = argmax.sequence_length as f32;
578                let char_positions = filtered_timesteps
579                    .iter()
580                    .map(|&timestep| timestep as f32 / seq_len)
581                    .collect();
582                let text = char_list.iter().collect();
583
584                (
585                    text,
586                    mean_conf,
587                    char_positions,
588                    filtered_timesteps,
589                    argmax.sequence_length,
590                )
591            })
592            .collect();
593
594        let mut all_texts = Vec::with_capacity(argmax.batch_size);
595        let mut all_scores = Vec::with_capacity(argmax.batch_size);
596        let mut all_positions = Vec::with_capacity(argmax.batch_size);
597        let mut all_col_indices = Vec::with_capacity(argmax.batch_size);
598        let mut all_seq_lengths = Vec::with_capacity(argmax.batch_size);
599        for (text, score, pos, cols, seq_len) in per {
600            all_texts.push(text);
601            all_scores.push(score);
602            all_positions.push(pos);
603            all_col_indices.push(cols);
604            all_seq_lengths.push(seq_len);
605        }
606
607        (
608            all_texts,
609            all_scores,
610            all_positions,
611            all_col_indices,
612            all_seq_lengths,
613        )
614    }
615
616    /// Applies the CTC decoder to a tensor of model predictions with character position tracking.
617    ///
618    /// This method handles the special requirements of CTC decoding and additionally tracks
619    /// the timestep positions of each character for word box generation.
620    ///
621    /// # Arguments
622    /// * `pred` - A 3D tensor containing the model predictions. Accepts any
623    ///   `ndarray` storage (owned `Array3<f32>` or a zero-copy `ArrayView3<f32>`).
624    ///
625    /// # Returns
626    /// A tuple containing:
627    /// * A vector of decoded text strings
628    /// * A vector of confidence scores for each text string
629    /// * A vector of character positions (normalized 0.0-1.0) for each text string
630    /// * A vector of column indices for each character in each text string
631    /// * A vector of sequence lengths (total columns) for each text string
632    pub fn apply_with_positions<S>(
633        &self,
634        pred: &ndarray::ArrayBase<S, ndarray::Ix3>,
635    ) -> PositionedDecodeResult
636    where
637        S: ndarray::Data<Elem = f32> + Sync,
638    {
639        if pred.is_empty() {
640            return (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new());
641        }
642        let argmax = self.argmax_predictions(pred);
643        self.decode_argmax_with_positions(&argmax)
644    }
645
646    /// Applies the CTC decoder to a tensor of model predictions.
647    ///
648    /// This method handles the special requirements of CTC decoding:
649    /// 1. Removing blank tokens
650    /// 2. Removing consecutive duplicate characters
651    /// 3. Converting indices to characters
652    /// 4. Calculating confidence scores
653    ///
654    /// # Arguments
655    /// * `pred` - A 3D tensor containing the model predictions. Accepts any
656    ///   `ndarray` storage (owned `Array3<f32>` or a zero-copy `ArrayView3<f32>`).
657    ///
658    /// # Returns
659    /// A tuple containing:
660    /// * A vector of decoded text strings
661    /// * A vector of confidence scores for each text string
662    pub fn apply<S>(&self, pred: &ndarray::ArrayBase<S, ndarray::Ix3>) -> (Vec<String>, Vec<f32>)
663    where
664        S: ndarray::Data<Elem = f32> + Sync,
665    {
666        if pred.is_empty() {
667            return (Vec::new(), Vec::new());
668        }
669        let argmax = self.argmax_predictions(pred);
670        self.decode_argmax(&argmax)
671    }
672}
673
674#[cfg(test)]
675mod tests {
676    use super::*;
677    use ndarray::Array3;
678
679    fn logits_with_winners(winners: &[&[(usize, f32)]], vocab_size: usize) -> Array3<f32> {
680        let batch_size = winners.len();
681        let sequence_length = winners.first().map_or(0, |sequence| sequence.len());
682        let mut logits = Array3::from_elem((batch_size, sequence_length, vocab_size), -10.0);
683        for (batch_idx, sequence) in winners.iter().enumerate() {
684            assert_eq!(sequence.len(), sequence_length);
685            for (time_idx, &(token_idx, probability)) in sequence.iter().enumerate() {
686                logits[[batch_idx, time_idx, token_idx]] = probability;
687            }
688        }
689        logits
690    }
691
692    #[test]
693    fn compact_argmax_preserves_ctc_text_scores_and_positions() {
694        let characters = vec!["a".to_string(), "b".to_string(), "c".to_string()];
695        let decoder = CTCLabelDecode::from_string_list(Some(&characters), false, false);
696        // Vocabulary slot 4 is deliberately out of the decoder's dictionary.
697        let logits = logits_with_winners(
698            &[
699                &[
700                    (0, 0.9),
701                    (1, 0.8),
702                    (1, 0.7),
703                    (0, 0.6),
704                    (1, 0.5),
705                    (2, 0.4),
706                    (2, 0.3),
707                ],
708                &[
709                    (3, 0.95),
710                    (3, 0.85),
711                    (4, 0.75),
712                    (3, 0.65),
713                    (0, 0.55),
714                    (2, 0.45),
715                    (0, 0.35),
716                ],
717            ],
718            5,
719        );
720
721        let argmax = decoder.argmax_predictions(&logits);
722        assert_eq!(argmax.batch_size, 2);
723        assert_eq!(argmax.sequence_length, 7);
724        assert_eq!(argmax.indices.len(), 14);
725        assert_eq!(argmax.probabilities.len(), 14);
726
727        let (texts, scores) = decoder.decode_argmax(&argmax);
728        assert_eq!(texts, ["aab", "ccb"]);
729        assert_eq!(
730            scores,
731            [(0.8 + 0.5 + 0.4) / 3.0, (0.95 + 0.65 + 0.45) / 3.0]
732        );
733
734        let (texts, scores, positions, columns, lengths) =
735            decoder.decode_argmax_with_positions(&argmax);
736        assert_eq!(texts, ["aab", "ccb"]);
737        assert_eq!(
738            scores,
739            [(0.8 + 0.5 + 0.4) / 3.0, (0.95 + 0.65 + 0.45) / 3.0]
740        );
741        assert_eq!(columns, [vec![1, 4, 5], vec![0, 3, 5]]);
742        assert_eq!(positions[0], [1.0 / 7.0, 4.0 / 7.0, 5.0 / 7.0]);
743        assert_eq!(positions[1], [0.0, 3.0 / 7.0, 5.0 / 7.0]);
744        assert_eq!(lengths, [7, 7]);
745    }
746
747    #[test]
748    fn compact_argmax_preserves_empty_tensor_behavior() {
749        let decoder = CTCLabelDecode::new(None, false);
750        let logits = Array3::<f32>::zeros((2, 0, decoder.get_character_count()));
751        let argmax = decoder.argmax_predictions(&logits);
752
753        assert_eq!(decoder.decode_argmax(&argmax), (Vec::new(), Vec::new()));
754        assert_eq!(
755            decoder.decode_argmax_with_positions(&argmax),
756            (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new())
757        );
758    }
759}