OneOf with bounding boxes crash

[Travvy88]

Hello! Thanks for good library for doc augmentations.

I faced problem with bounding boxes. When I run augraphy pipeline with bounding boxes and there is the OneOf inside it crashes. When I delete OneOf everything works well.

import random
import numpy as np
from src.augmentations import get_augmentation_phases
import augraphy
from PIL import Image, ImageDraw

ink_phase_one_of = [
    augraphy.OneOf(
            [
                augraphy.Dithering(
                    dither=random.choice(["ordered", "floyd-steinberg"]),
                    order=(3, 5),
                ),
                augraphy.InkBleed(
                    intensity_range=(0.1, 0.2),
                    kernel_size=random.choice([(7, 7), (5, 5), (3, 3)]),
                    severity=(0.4, 0.6),
                ),
            ],
            p=1,
        ),
]

ink_phase = [
    augraphy.Dithering(
        dither=random.choice(["ordered", "floyd-steinberg"]),
        order=(3, 5),
    ),
    augraphy.InkBleed(
        intensity_range=(0.1, 0.2),
        kernel_size=random.choice([(7, 7), (5, 5), (3, 3)]),
        severity=(0.4, 0.6),
    ),
]

image  = Image.open('data/im_2.png')

bounding_boxes = [
    [150, 90, 235, 105],
    [150, 150, 298, 170]
]
pipeline = augraphy.AugraphyPipeline(
                                     bounding_boxes=bounding_boxes,
                                     log=True, 
                                     paper_phase=[],
                                     ink_phase=ink_phase_one_of,
                                     post_phase=[],
                                     pre_phase=[])

output = pipeline.augment(np.array(image))

Error log:

Traceback (most recent call last):
  File "/home/dedoc/shevtsov/WordPaletteDocumenter/test_aug.py", line 50, in <module>
    output = pipeline.augment(np.array(image))
  File "/home/dedoc/.virtualenvs/wpd/lib/python3.10/site-packages/augraphy/base/augmentationpipeline.py", line 213, in augment
    data = self.augment_single_image(
  File "/home/dedoc/.virtualenvs/wpd/lib/python3.10/site-packages/augraphy/base/augmentationpipeline.py", line 375, in augment_single_image
    self.apply_phase(data, layer="ink", phase=self.ink_phase)
  File "/home/dedoc/.virtualenvs/wpd/lib/python3.10/site-packages/augraphy/base/augmentationpipeline.py", line 758, in apply_phase
    result, mask, keypoints, bounding_boxes = result
ValueError: not enough values to unpack (expected 4, got 2)

It seems like augmentationpipeline.py:754 line should filter OneOf and AugmentationSequence objects. However, both of them are children of Augmentation class. isinstance(augmentation, Augmentation) is True when augmentation variable is OneOf or AugmentationSequence and this objects are not filtered.

[Travvy88]

I propose to change the if statement:

if isinstance(augmentation, Augmentation) and  \
                        not isinstance(augmentation, OneOf) and \
                        not isinstance(augmentation, AugmentationSequence):

Everything works okay until new error is thrown:

Traceback (most recent call last):
  File "/home/dedoc/shevtsov/WordPaletteDocumenter/test_aug.py", line 48, in <module>
    output = pipeline.augment(np.array(image))
  File "/home/dedoc/.virtualenvs/wpd/lib/python3.10/site-packages/augraphy/base/augmentationpipeline.py", line 214, in augment
    data = self.augment_single_image(
  File "/home/dedoc/.virtualenvs/wpd/lib/python3.10/site-packages/augraphy/base/augmentationpipeline.py", line 402, in augment_single_image
    self.apply_phase(data, layer="paper", phase=self.paper_phase)
  File "/home/dedoc/.virtualenvs/wpd/lib/python3.10/site-packages/augraphy/base/augmentationpipeline.py", line 798, in apply_phase
    result, mask, keypoints, bounding_boxes = result
ValueError: not enough values to unpack (expected 4, got 2)

Seems like 798 line should process OneOf and AugmentationSequence objects. But there are no bboxes, masks and keypoints in result variable and it crashes every time.

[jboarman]

@Travvy88 Do you mind investigating this further to see if you can identify a solution? If so, you could prepare a PR and we could include you as a listed contributor associated with the library.

[Travvy88]

@Travvy88 Do you mind investigating this further to see if you can identify a solution? If so, you could prepare a PR and we could include you as a listed contributor associated with the library.

Yes, I will try.

[Travvy88]

@jboarman according to my research, the next error at line augmentationpipeline.py:791 is caused when the program tries to extract bboxes, keypoints and mask from result after applying AugmentationSequence. However, AugmentationSequence return only image (current_result) and list of augmentations:

def __call__(self, image, layer=None, mask=None, keypoints=None, bounding_boxes=None, force=False):
        if force or self.should_run():
            # reset to prevent memory leaks
            self.results = []
            result = image
            for augmentation in self.augmentations:
                if isinstance(result, tuple):
                    result = result[0]
                current_result = augmentation(result, mask=mask, keypoints=keypoints, bounding_boxes=bounding_boxes)

                if isinstance(augmentation, Augmentation):
                    if (mask is not None) or (keypoints is not None) or (bounding_boxes is not None):
                        current_result, mask, keypoints, bounding_boxes = current_result
                self.results.append(current_result)

                # make sure result is not None when parsing it to the next augmentation
                if not isinstance(result, tuple) and current_result is not None:
                    result = current_result
                elif isinstance(current_result, tuple):
                    if current_result[0] is not None:
                        result = current_result
                        
            return result, self.augmentations

Can you tell me, is that the plan? Or there is possibly a bug and the bboxes and other info should be nested into result variable to return it from AugmentationSequence?

[kwcckw]

Thanks for pointing this out. In augmentationpipeline.py, line 752, for OneOf, we can fix it with:

            if isinstance(augmentation, Augmentation):
                # "OneOf" 
                if augmentation.__class__.__name__ == "OneOf":
                    result, augmentations = result
                    if (mask is not None) or (keypoints is not None) or (bounding_boxes is not None):
                        result, mask, keypoints, bounding_boxes = result
                # "AugmentationSequence"
                elif augmentation.__class__.__name__ == "AugmentationSequence":
                    pass
                # others
                else:
                    # unpacking augmented image, mask, keypoints and bounding boxes from output
                    if (mask is not None) or (keypoints is not None) or (bounding_boxes is not None):
                        result, mask, keypoints, bounding_boxes = result

Could you check and do the same for AugmentationSequence too?

[Travvy88]

I fixed 751st line in apply_phase and some code inside OneOf and AugmentationSequence. It seems that these changes are enough to solve the problem. You can check my PR #446.