Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Updated TextImage #1860

Merged
merged 1 commit into from
Jul 30, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions albumentations/augmentations/text/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
ternaus marked this conversation as resolved.
Show resolved Hide resolved
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)
Expand All @@ -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,
ternaus marked this conversation as resolved.
Show resolved Hide resolved
"text": augmented_text,
"font": font,
"font_color": font_color,
Expand All @@ -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 {
Expand All @@ -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"]
Expand Down
Loading