Updated TextImage (#1860)

albumentations-team · Jul 30, 2024 · 512dc63 · 512dc63
1 parent 33c8927
commit 512dc63
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/albumentations/augmentations/text/transforms.py b/albumentations/augmentations/text/transforms.py
@@ -154,7 +154,7 @@ def random_aug(
         result_sentence = re.sub(" +", " ", result_sentence).strip()
         return result_sentence if result_sentence != text else ""
 
-    def preprocess_metadata(self, image: np.ndarray, bbox: BoxType, text: str) -> dict[str, Any]:
+    def preprocess_metadata(self, image: np.ndarray, bbox: BoxType, text: str, bbox_index: int) -> dict[str, Any]:
         image_height, image_width = image.shape[:2]
 
         check_bbox(bbox)
@@ -178,6 +178,8 @@ def preprocess_metadata(self, image: np.ndarray, bbox: BoxType, text: str) -> di
 
         return {
             "bbox_coords": (x_min, y_min, x_max, y_max),
+            "bbox_index": bbox_index,
+            "original_text": text,
             "text": augmented_text,
             "font": font,
             "font_color": font_color,
@@ -203,8 +205,8 @@ def get_params_dependent_on_data(self, params: dict[str, Any], data: dict[str, A
         bbox_indices_to_update = random.sample(range(len(metadata)), num_bboxes_to_modify)
 
         overlay_data = [
-            self.preprocess_metadata(image, metadata[index]["bbox"], metadata[index]["text"])
-            for index in bbox_indices_to_update
+            self.preprocess_metadata(image, metadata[bbox_index]["bbox"], metadata[bbox_index]["text"], bbox_index)
+            for bbox_index in bbox_indices_to_update
         ]
 
         return {
@@ -225,6 +227,8 @@ def apply_with_params(self, params: dict[str, Any], *args: Any, **kwargs: Any) -
             {
                 "bbox_coords": overlay["bbox_coords"],
                 "text": overlay["text"],
+                "original_text": overlay["original_text"],
+                "bbox_index": overlay["bbox_index"],
                 "font_color": overlay["font_color"],
             }
             for overlay in params["overlay_data"]