Bump Python requirement torchvision to 0.21.0

dan/ocr/predict/attention.py

@@ -2,7 +2,6 @@
 import logging
 import re
 from enum import Enum
-from itertools import pairwise
 from typing import Dict, List, Tuple
 
 import cv2
@@ -28,78 +27,30 @@ class Level(str, Enum):
 
 
 def parse_delimiters(delimiters: List[str]) -> re.Pattern:
-    return re.compile(r"|".join(delimiters))
+    return re.compile(rf"[^{'|'.join(delimiters)}]+")
 
 
-def build_ner_indices(
-    text: str, tokens: Dict[str, EntityType]
-) -> List[Tuple[int, int]]:
-    """
-    Compute the position of NER tokens in the text and return a list of indices.
-    :param text: list of characters.
-    :param tokens: NER tokens used.
-    Returns a list of indices where tokens are located.
-    """
-    start_tokens, end_tokens = zip(*list(tokens.values()))
-    end_tokens = list(filter(bool, end_tokens))
-
-    if len(end_tokens):
-        assert len(start_tokens) == len(
-            end_tokens
-        ), "You don't have the same number of starting tokens and ending tokens"
-        return [
-            [pos_start, pos_end] for pos_start, pos_end in zip(start_tokens, end_tokens)
-        ]
-
-    return list(
-        pairwise(
-            [pos for pos, char in enumerate(text) if char in start_tokens] + [None]
-        )
-    )
-
-
-def compute_offsets_by_level(
-    level: Level, text_list: List[str], indices: List[Tuple[int, int]]
-):
+def compute_offsets_by_level(full_text: str, level: Level, text_list: List[str]):
     """
     Compute and return the list of offset between each text part.
+    :param full_text: predicted text.
     :param level: Level to use from [char, word, line, ner].
     :param text_list: list of text to use.
-    :param indices: list of indices where tokens are located for NER computation.
     Returns a list of offsets.
     """
-    if level == Level.NER:
-        return (
-            [
-                current - next_token
-                for (_, next_token), (current, _) in pairwise(indices)
-            ]
-            # Pad the list to match the length of the text list
-            + [0]
-        )
-
-    return [int(level != Level.Char)] * len(text_list)
+    # offsets[idx] = number of characters between text_list[idx-1] and text_list[idx]
+    offsets = [int(level != Level.Char)] * (len(text_list) - 1)
 
+    if level == Level.NER:
+        # Start after the first entity
+        cursor = len(text_list[0])
+        for idx, split in enumerate(text_list[1:]):
+            # Number of characters between this entity and the previous one
+            offsets[idx] = full_text[cursor:].index(split)
+            cursor += offsets[idx] + len(split)
 
-def compute_prob_by_ner(
-    characters: str, probabilities: List[float], indices: List[Tuple[int, int]]
-) -> Tuple[List[str], List[np.float64]]:
-    """
-    Split text and confidences using indices and return a list of average confidence scores.
-    :param characters: list of characters.
-    :param probabilities: list of character probabilities.
-    :param indices: list of indices where tokens are located.
-    Returns a list confidence scores.
-    """
-    return zip(
-        *[
-            (
-                characters[current:next_token],
-                np.mean(probabilities[current:next_token]),
-            )
-            for current, next_token in indices
-        ]
-    )
+    # Last offset is not used, padded with a 0 to match the length of text_list
+    return offsets + [0]
 
 
 def compute_prob_by_separator(
@@ -113,78 +64,75 @@ def compute_prob_by_separator(
     Returns a list confidence scores.
     """
     # match anything except separators, get start and end index
-    pattern = re.compile(f"[^{separator.pattern}]+")
-    matches = [(m.start(), m.end()) for m in re.finditer(pattern, characters)]
+    matches = [(m.start(), m.end()) for m in separator.finditer(characters)]
 
     # Iterate over text pieces and compute mean confidence
-    probs = [np.mean(probabilities[start:end]) for (start, end) in matches]
-    texts = [characters[start:end] for (start, end) in matches]
-    return texts, probs
+    return [characters[start:end] for (start, end) in matches], [
+        np.mean(probabilities[start:end]) for (start, end) in matches
+    ]
 
 
 def split_text(
     text: str,
     level: Level,
+    char_separators: re.Pattern,
     word_separators: re.Pattern,
     line_separators: re.Pattern,
-    tokens: Dict[str, EntityType],
+    tokens_separators: re.Pattern | None = None,
 ) -> Tuple[List[str], List[int]]:
     """
     Split text into a list of characters, word, or lines.
     :param text: Text prediction from DAN
     :param level: Level to visualize from [char, word, line, ner]
-    :param word_separators: List of word separators
-    :param line_separators: List of line separators
-    :param tokens: NER tokens used
+    :param char_separators: Pattern used to find tokens in the charset
+    :param word_separators: Pattern used to find words
+    :param line_separators: Pattern used to find lines
+    :param tokens_separators: Pattern used to find NER entities
     """
-    indices = []
-
     match level:
         case Level.Char:
-            text_split = list(text)
+            text_split = char_separators.findall(text)
         # split into words
         case Level.Word:
-            text_split = re.split(word_separators, text)
+            text_split = word_separators.split(text)
         # split into lines
         case Level.Line:
-            text_split = re.split(line_separators, text)
+            text_split = line_separators.split(text)
         # split into entities
         case Level.NER:
-            if not tokens:
+            if not tokens_separators:
                 logger.error("Cannot compute NER level: tokens not found")
                 return [], []
-
-            indices = build_ner_indices(text, tokens)
-            text_split = [text[current:next_token] for current, next_token in indices]
+            text_split = tokens_separators.findall(text)
         case _:
             logger.error(f"Level should be either {list(map(str, Level))}")
             return [], []
 
-    return text_split, compute_offsets_by_level(level, text_split, indices)
+    return text_split, compute_offsets_by_level(text, level, text_split)
 
 
 def split_text_and_confidences(
     text: str,
     confidences: List[float],
     level: Level,
+    char_separators: re.Pattern,
     word_separators: re.Pattern,
     line_separators: re.Pattern,
-    tokens: Dict[str, EntityType],
+    tokens_separators: re.Pattern | None = None,
 ) -> Tuple[List[str], List[np.float64], List[int]]:
     """
     Split text into a list of characters, words or lines with corresponding confidences scores
     :param text: Text prediction from DAN
     :param confidences: Character confidences
     :param level: Level to visualize from [char, word, line, ner]
-    :param word_separators: List of word separators
-    :param line_separators: List of line separators
-    :param tokens: NER tokens used
+    :param char_separators: Pattern used to find tokens of the charset
+    :param word_separators: Pattern used to find words
+    :param line_separators: Pattern used to find lines
+    :param tokens_separators: Pattern used to find NER entities
     """
-    indices = []
-
     match level:
         case Level.Char:
-            texts = list(text)
+            texts = char_separators.findall(text)
         case Level.Word:
             texts, confidences = compute_prob_by_separator(
                 text, confidences, word_separators
@@ -194,15 +142,13 @@ def split_text_and_confidences(
                 text, confidences, line_separators
             )
         case Level.NER:
-            if not tokens:
+            if not tokens_separators:
                 logger.error("Cannot compute NER level: tokens not found")
                 return [], [], []
 
-            indices = build_ner_indices(text, tokens)
-            if not indices:
-                return [], [], []
-
-            texts, confidences = compute_prob_by_ner(text, confidences, indices)
+            texts, confidences = compute_prob_by_separator(
+                text, confidences, tokens_separators
+            )
         case _:
             logger.error(f"Level should be either {list(map(str, Level))}")
             return [], [], []
@@ -210,7 +156,7 @@ def split_text_and_confidences(
     return (
         texts,
         [np.around(num, 2) for num in confidences],
-        compute_offsets_by_level(level, texts, indices),
+        compute_offsets_by_level(text, level, texts),
     )
 
 
@@ -221,10 +167,11 @@ def get_predicted_polygons_with_confidence(
     level: Level,
     height: int,
     width: int,
+    char_separators: re.Pattern,
     max_object_height: int = 50,
     word_separators: re.Pattern = parse_delimiters(["\n", " "]),
     line_separators: re.Pattern = parse_delimiters(["\n"]),
-    tokens: Dict[str, EntityType] = {},
+    tokens_separators: re.Pattern | None = None,
 ) -> List[dict]:
     """
     Returns the polygons of each object of the current prediction
@@ -234,14 +181,21 @@ def get_predicted_polygons_with_confidence(
     :param level: Level to display (must be in [char, word, line, ner])
     :param height: Original image height
     :param width: Original image width
+    :param char_separators: Pattern used to find tokens of the charset
     :param max_object_height: Maximum height of predicted objects.
-    :param word_separators: List of word separators
-    :param line_separators: List of line separators
-    :param tokens: NER tokens used
+    :param word_separators: Pattern used to find words
+    :param line_separators: Pattern used to find lines
+    :param tokens_separators: Pattern used to find NER entities
     """
     # Split text into characters, words or lines
     text_list, confidence_list, offsets = split_text_and_confidences(
-        text, confidences, level, word_separators, line_separators, tokens
+        text,
+        confidences,
+        level,
+        char_separators,
+        word_separators,
+        line_separators,
+        tokens_separators,
     )
 
     max_value = weights.sum(0).max()