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.

Sign up for GitHub

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

[Fix] Remove image padding after rotation correction with straighten_pages=True #1731

Merged
merged 2 commits into from
Sep 30, 2024
Merged

[Fix] Remove image padding after rotation correction with straighten_pages=True #1731

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ae358ea
add remove image padding function
felixdittrich92 Sep 25, 2024
489a7c4
Update doctr/models/predictor/base.py
felixdittrich92 Sep 30, 2024
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
6 changes: 3 additions & 3 deletions doctr/models/predictor/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import numpy as np

from doctr.models.builder import DocumentBuilder
from doctr.utils.geometry import extract_crops, extract_rcrops, rotate_image
from doctr.utils.geometry import extract_crops, extract_rcrops, remove_image_padding, rotate_image

from .._utils import estimate_orientation, rectify_crops, rectify_loc_preds
from ..classification import crop_orientation_predictor, page_orientation_predictor
Expand Down Expand Up @@ -101,8 +101,8 @@ def _straighten_pages(
]
)
return [
# expand if height and width are not equal
rotate_image(page, angle, expand=page.shape[0] != page.shape[1])
# expand if height and width are not equal, then remove the padding
remove_image_padding(rotate_image(page, angle, expand=page.shape[0] != page.shape[1]))
for page, angle in zip(pages, origin_pages_orientations)
]

Expand Down
21 changes: 21 additions & 0 deletions doctr/utils/geometry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
"rotate_boxes",
"compute_expanded_shape",
"rotate_image",
"remove_image_padding",
"estimate_page_angle",
"convert_to_relative_coords",
"rotate_abs_geoms",
Expand Down Expand Up @@ -351,6 +352,26 @@ def rotate_image(
return rot_img


def remove_image_padding(image: np.ndarray) -> np.ndarray:
"""Remove black border padding from an image

Args:
----
image: numpy tensor to remove padding from

Returns:
-------
Image with padding removed
"""
# Find the bounding box of the non-black region
rows = np.any(image, axis=1)
cols = np.any(image, axis=0)
rmin, rmax = np.where(rows)[0][[0, -1]]
cmin, cmax = np.where(cols)[0][[0, -1]]

return image[rmin : rmax + 1, cmin : cmax + 1]


def estimate_page_angle(polys: np.ndarray) -> float:
"""Takes a batch of rotated previously ORIENTED polys (N, 4, 2) (rectified by the classifier) and return the
estimated angle ccw in degrees
Expand Down
11 changes: 11 additions & 0 deletions tests/common/test_utils_geometry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,17 @@ def test_rotate_image():
assert rotated[0, :, 0].sum() <= 1


def test_remove_image_padding():
img = np.ones((32, 64, 3), dtype=np.float32)
padded = np.pad(img, ((10, 10), (20, 20), (0, 0)))
cropped = geometry.remove_image_padding(padded)
assert np.all(cropped == img)

# No padding
cropped = geometry.remove_image_padding(img)
assert np.all(cropped == img)


@pytest.mark.parametrize(
"abs_geoms, img_size, rel_geoms",
[
Expand Down
Loading