From b8f12ea459b096787552421d1abf7045c24a8b1a Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 7 Jul 2026 10:04:13 +0200 Subject: [PATCH] [misc] Apply safe guards and minor improvements --- doctr/io/image/base.py | 5 +++-- doctr/io/image/pytorch.py | 8 +++++--- doctr/models/builder.py | 2 ++ doctr/models/detection/core.py | 3 +-- doctr/models/detection/fast/base.py | 2 +- doctr/models/detection/linknet/base.py | 2 +- doctr/models/predictor/pytorch.py | 2 ++ tests/common/test_io.py | 15 +++++++++++++++ tests/common/test_models_builder.py | 19 +++++++++++++++++++ tests/pytorch/test_io_image_pt.py | 17 +++++++++++++++++ 10 files changed, 66 insertions(+), 9 deletions(-) diff --git a/doctr/io/image/base.py b/doctr/io/image/base.py index 0c75b1b72b..69c3aa137a 100644 --- a/doctr/io/image/base.py +++ b/doctr/io/image/base.py @@ -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'") diff --git a/doctr/io/image/pytorch.py b/doctr/io/image/pytorch.py index cc32cb2543..50a60be491 100644 --- a/doctr/io/image/pytorch.py +++ b/doctr/io/image/pytorch.py @@ -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 @@ -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: @@ -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: diff --git a/doctr/models/builder.py b/doctr/models/builder.py index 23c42726d6..149fc3c7d8 100644 --- a/doctr/models/builder.py +++ b/doctr/models/builder.py @@ -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 diff --git a/doctr/models/detection/core.py b/doctr/models/detection/core.py index d7edfe9367..a3d43e03ec 100644 --- a/doctr/models/detection/core.py +++ b/doctr/models/detection/core.py @@ -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, diff --git a/doctr/models/detection/fast/base.py b/doctr/models/detection/fast/base.py index 50a9ae17c0..1319367e86 100644 --- a/doctr/models/detection/fast/base.py +++ b/doctr/models/detection/fast/base.py @@ -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: diff --git a/doctr/models/detection/linknet/base.py b/doctr/models/detection/linknet/base.py index db0f18e5ed..db799f0149 100644 --- a/doctr/models/detection/linknet/base.py +++ b/doctr/models/detection/linknet/base.py @@ -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: diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index 56377c555b..a7c36d803a 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -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"], diff --git a/tests/common/test_io.py b/tests/common/test_io.py index 8e5fd1118d..fce0abaa41 100644 --- a/tests/common/test_io.py +++ b/tests/common/test_io.py @@ -4,6 +4,7 @@ import numpy as np import pytest import requests +from PIL import Image from doctr import io @@ -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) diff --git a/tests/common/test_models_builder.py b/tests/common/test_models_builder.py index 270701452f..2c1084b8a1 100644 --- a/tests/common/test_models_builder.py +++ b/tests/common/test_models_builder.py @@ -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"] diff --git a/tests/pytorch/test_io_image_pt.py b/tests/pytorch/test_io_image_pt.py index 2c1ab69c0b..a3867d01b9 100644 --- a/tests/pytorch/test_io_image_pt.py +++ b/tests/pytorch/test_io_image_pt.py @@ -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 @@ -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)