[misc] minor performance & memory improvements

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:

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)
 

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
@@ -223,23 +225,12 @@ 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)
     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

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():