Skip to content
Merged
Show file tree
Hide file tree
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
5 changes: 3 additions & 2 deletions doctr/io/image/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,10 @@ def read_img_as_numpy(
if isinstance(file, (str, Path)):
if not Path(file).is_file():
raise FileNotFoundError(f"unable to access {file}")
img = cv2.imread(str(file), cv2.IMREAD_COLOR)
_file: np.ndarray = np.frombuffer(Path(file).read_bytes(), np.uint8)
img = cv2.imdecode(_file, cv2.IMREAD_COLOR)
elif isinstance(file, bytes):
_file: np.ndarray = np.frombuffer(file, np.uint8)
_file = np.frombuffer(file, np.uint8)
img = cv2.imdecode(_file, cv2.IMREAD_COLOR)
else:
raise TypeError("unsupported object type for argument 'file'")
Expand Down
8 changes: 5 additions & 3 deletions doctr/io/image/pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

import numpy as np
import torch
from PIL import Image
from PIL import Image, ImageOps
from torchvision.transforms.functional import to_tensor

from doctr.utils.common_types import AbstractPath
Expand Down Expand Up @@ -47,7 +47,9 @@ def read_img_as_tensor(img_path: AbstractPath, dtype: torch.dtype = torch.float3
raise ValueError("insupported value for dtype")

with Image.open(img_path, mode="r") as pil_img:
return tensor_from_pil(pil_img.convert("RGB"), dtype)
# Apply the EXIF orientation before decoding, to stay consistent with `read_img_as_numpy`
# (OpenCV applies EXIF orientation automatically, PIL does not)
return tensor_from_pil(ImageOps.exif_transpose(pil_img).convert("RGB"), dtype)


def decode_img_as_tensor(img_content: bytes, dtype: torch.dtype = torch.float32) -> torch.Tensor:
Expand All @@ -64,7 +66,7 @@ def decode_img_as_tensor(img_content: bytes, dtype: torch.dtype = torch.float32)
raise ValueError("insupported value for dtype")

with Image.open(BytesIO(img_content), mode="r") as pil_img:
return tensor_from_pil(pil_img.convert("RGB"), dtype)
return tensor_from_pil(ImageOps.exif_transpose(pil_img).convert("RGB"), dtype)


def tensor_from_numpy(npy_img: np.ndarray, dtype: torch.dtype = torch.float32) -> torch.Tensor:
Expand Down
2 changes: 2 additions & 0 deletions doctr/models/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -386,6 +386,8 @@ def _build_tables(
tables_out: list[Table] = []
for table_dict in table_dicts:
cells = table_dict["cells"]
if len(cells) == 0:
continue
cell_polys = [self._as_cell_polygon(cell["geometry"]) for cell in cells]

# Assign each (still unassigned) word to at most one cell of this table: the first cell (in cell
Expand Down
3 changes: 1 addition & 2 deletions doctr/models/detection/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,7 @@ def box_score(pred: np.ndarray, points: np.ndarray, assume_straight_pages: bool
mask: np.ndarray = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
cv2.fillPoly(mask, [pts - np.array([[xmin, ymin]], dtype=np.int32)], 1)
vals = pred[ymin : ymax + 1, xmin : xmax + 1][mask.astype(bool)]
nonzero = np.count_nonzero(vals)
return float(vals.sum() / nonzero) if nonzero > 0 else 0.0
return float(vals.mean()) if vals.size > 0 else 0.0

def bitmap_to_boxes(
self,
Expand Down
2 changes: 1 addition & 1 deletion doctr/models/detection/fast/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ def bitmap_to_boxes(
else:
_box = self.polygon_to_box(np.squeeze(contour))

if _box is None:
if _box is None: # pragma: no cover
continue

if self.assume_straight_pages:
Expand Down
2 changes: 1 addition & 1 deletion doctr/models/detection/linknet/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ def bitmap_to_boxes(
else:
_box = self.polygon_to_box(np.squeeze(contour))

if _box is None:
if _box is None: # pragma: no cover
continue

if self.assume_straight_pages:
Expand Down
2 changes: 2 additions & 0 deletions doctr/models/predictor/pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,8 @@ def _tables_from_regions(
new_cell = dict(cell)
new_cell["geometry"] = poly.tolist()
remapped_cells.append(new_cell)
if not remapped_cells:
continue
tables_per_page[p_idx].append({
"cells": remapped_cells,
"num_rows": grid["num_rows"],
Expand Down
15 changes: 15 additions & 0 deletions tests/common/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import numpy as np
import pytest
import requests
from PIL import Image

from doctr import io

Expand Down Expand Up @@ -97,3 +98,17 @@ def test_pdf(mock_pdf):
# As images
num_pages = 2
_check_doc_content(pages, num_pages)


def test_read_img_as_numpy_exif_orientation(tmpdir_factory):
# A JPEG with EXIF orientation 6 (90° clockwise display rotation)
folder = tmpdir_factory.mktemp("images")
path = str(folder.join("exif_o6.jpg"))
pil_img = Image.fromarray(np.zeros((40, 60, 3), dtype=np.uint8))
exif = pil_img.getexif()
exif[0x0112] = 6 # orientation tag
pil_img.save(path, format="JPEG", exif=exif)

assert io.read_img_as_numpy(path).shape == (60, 40, 3)
with open(path, "rb") as f:
assert io.read_img_as_numpy(f.read()).shape == (60, 40, 3)
19 changes: 19 additions & 0 deletions tests/common/test_models_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -562,3 +562,22 @@ def test_sort_boxes_degenerate_heights():
boxes = np.array([[0.5, 0.2, 0.6, 0.2], [0.1, 0.2, 0.2, 0.2]], dtype=np.float32)
idxs, _ = doc_builder._sort_boxes(boxes)
assert sorted(np.asarray(idxs).tolist()) == [0, 1]


def test_documentbuilder_tables_empty_cells():
# A table prediction with no cells (e.g. a false-positive "Table" region where the table model
# finds nothing) must not crash the document build
doc_builder = builder.DocumentBuilder()
boxes = np.array([[0.1, 0.1, 0.2, 0.2]], dtype=np.float32)
out = doc_builder(
[np.zeros((100, 100, 3))],
[boxes],
[np.array([0.9])],
[[("hello", 0.99)]],
[(100, 100)],
[[{"value": 0, "confidence": None}]],
tables=[[{"cells": [], "num_rows": 0, "num_cols": 0}]],
)
assert out.pages[0].tables == []
# the word stays in the regular blocks since it was never consumed by a table
assert [w.value for b in out.pages[0].blocks for line in b.lines for w in line.words] == ["hello"]
17 changes: 17 additions & 0 deletions tests/pytorch/test_io_image_pt.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import numpy as np
import pytest
import torch
from PIL import Image

from doctr.io import decode_img_as_tensor, read_img_as_tensor, tensor_from_numpy

Expand Down Expand Up @@ -51,3 +52,19 @@ def test_tensor_from_numpy(mock_image_stream):
assert out.dtype == torch.float16
out = tensor_from_numpy(np.zeros((256, 256, 3), dtype=np.uint8), dtype=torch.uint8)
assert out.dtype == torch.uint8


def test_read_img_as_tensor_exif_orientation(tmpdir_factory):
folder = tmpdir_factory.mktemp("images")
path = str(folder.join("exif_o6.jpg"))
pil_img = Image.fromarray(np.zeros((40, 60, 3), dtype=np.uint8))
exif = pil_img.getexif()
exif[0x0112] = 6 # orientation tag
pil_img.save(path, format="JPEG", exif=exif)

img = read_img_as_tensor(path)
assert img.shape == (3, 60, 40)

with open(path, "rb") as f:
img_stream = decode_img_as_tensor(f.read())
assert img_stream.shape == (3, 60, 40)
Loading