From 5ed42f7f7cc041ea4b9727ea7a0cc44b331b8efd Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 28 Jul 2026 11:55:54 +0200 Subject: [PATCH 1/2] [misc] minor performance & memory improvements --- doctr/models/kie_predictor/pytorch.py | 2 +- doctr/models/predictor/pytorch.py | 2 +- doctr/models/reading_order/base.py | 8 +++----- tests/pytorch/test_models_detection_pt.py | 10 +++++----- 4 files changed, 10 insertions(+), 12 deletions(-) diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index ba84bd8523..6e71d4e5c0 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/doctr/models/kie_predictor/pytorch.py @@ -97,7 +97,7 @@ def forward( loc_preds, out_maps = self.det_predictor(det_pages, return_maps=True, **kwargs) bin_thresh = kwargs.get("bin_thresh", getattr(self.det_predictor.model.postprocessor, "bin_thresh", 0.3)) seg_maps = [ - ((np.expand_dims(np.amax(out_map, axis=-1), axis=-1) > bin_thresh) * 255).astype(np.uint8) + ((np.expand_dims(np.amax(out_map, axis=-1), axis=-1) > bin_thresh).astype(np.uint8) * 255) for out_map in out_maps ] else: diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index 41dc228320..cf76d15ab8 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -110,7 +110,7 @@ def forward( if self.detect_orientation or self.straighten_pages: loc_preds, out_maps = self.det_predictor(det_pages, return_maps=True, **kwargs) bin_thresh = getattr(self.det_predictor.model.postprocessor, "bin_thresh") - seg_maps = [((out_map > bin_thresh) * 255).astype(np.uint8) for out_map in out_maps] + seg_maps = [((out_map > bin_thresh).astype(np.uint8) * 255) for out_map in out_maps] else: loc_preds = self.det_predictor(det_pages, **kwargs) diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index bafba48a27..b1fa463386 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -8,6 +8,8 @@ from typing import Any import numpy as np +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import connected_components from doctr.utils.geometry import estimate_page_angle, order_points from doctr.utils.repr import NestedObject @@ -235,11 +237,7 @@ def _find(node: int) -> int: col_edges = np.triu(x_linked, 1) col_edges &= ~spanning[:, None] col_edges &= ~spanning[None, :] - for i, j in np.argwhere(col_edges): - ri, rj = _find(int(i)), _find(int(j)) - if ri != rj: - parent[ri] = rj - component = np.array([_find(i) for i in range(num_boxes)]) + component = connected_components(csr_matrix(col_edges), directed=False, return_labels=True)[1] # Detect whether the page is multi-column: if a vertical line can be drawn that separates the boxes into two # groups with a small number of crossing boxes, the page is considered multi-column. This is used to diff --git a/tests/pytorch/test_models_detection_pt.py b/tests/pytorch/test_models_detection_pt.py index 40679a53ef..d29513f4b0 100644 --- a/tests/pytorch/test_models_detection_pt.py +++ b/tests/pytorch/test_models_detection_pt.py @@ -108,13 +108,13 @@ def test_detection_zoo(arch_name): predictor.model.cuda() with torch.no_grad(): - out, seq_maps = predictor(input_tensor, return_maps=True) + out, seg_maps = predictor(input_tensor, return_maps=True) assert all(isinstance(boxes, dict) for boxes in out) assert all(isinstance(boxes[CLASS_NAME], np.ndarray) and boxes[CLASS_NAME].shape[1] == 5 for boxes in out) - assert all(isinstance(seq_map, np.ndarray) for seq_map in seq_maps) - assert all(seq_map.shape[:2] == (1024, 1024) for seq_map in seq_maps) - # check that all values in the seq_maps are between 0 and 1 - assert all((seq_map >= 0).all() and (seq_map <= 1).all() for seq_map in seq_maps) + assert all(isinstance(seg_map, np.ndarray) for seg_map in seg_maps) + assert all(seg_map.shape[:2] == (1024, 1024) for seg_map in seg_maps) + # check that all values in the seg_maps are between 0 and 1 + assert all((seg_map >= 0).all() and (seg_map <= 1).all() for seg_map in seg_maps) def test_fast_reparameterization(): From 2e9c5b7b262d1574385e25f016eee0c9785269e1 Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 28 Jul 2026 11:59:20 +0200 Subject: [PATCH 2/2] Remove unused code --- doctr/models/reading_order/base.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/doctr/models/reading_order/base.py b/doctr/models/reading_order/base.py index b1fa463386..923db2b1f5 100644 --- a/doctr/models/reading_order/base.py +++ b/doctr/models/reading_order/base.py @@ -225,13 +225,6 @@ def _topological_order(boxes: np.ndarray, x_overlap_threshold: float, y_overlap_ # the direct top-to-bottom continuation is momentarily broken by a fragmented or misaligned OCR line. page_width = float(x1.max() - x0.min()) or 1.0 spanning = (x1 - x0) > 0.5 * page_width - parent = np.arange(num_boxes) - - def _find(node: int) -> int: - while parent[node] != node: - parent[node] = parent[parent[node]] - node = int(parent[node]) - return node # Reuse the horizontal-overlap matrix; keep the upper triangle only, so each pair is visited once col_edges = np.triu(x_linked, 1)