[misc] Apply safe guards and minor improvements (#2104)

Author: felixdittrich92

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'")

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:

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

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,

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:

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:

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"],

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)

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"]

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)