From d93173b072823dbcd4ba3b7c3fed8281af5cfa5d Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 3 Jul 2026 15:55:19 +0200 Subject: [PATCH] rotated post proc improvements --- doctr/models/detection/core.py | 14 ++++++++++---- doctr/models/preprocessor/pytorch.py | 1 - 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/doctr/models/detection/core.py b/doctr/models/detection/core.py index 9c8c7bc1a0..d7edfe9367 100644 --- a/doctr/models/detection/core.py +++ b/doctr/models/detection/core.py @@ -52,10 +52,16 @@ def box_score(pred: np.ndarray, points: np.ndarray, assume_straight_pages: bool return pred[ymin : ymax + 1, xmin : xmax + 1].mean() else: - mask: np.ndarray = np.zeros((h, w), np.int32) - cv2.fillPoly(mask, [points.astype(np.int32)], 1.0) - product = pred * mask - return np.sum(product) / np.count_nonzero(product) + pts: np.ndarray = points.reshape((-1, 2)).astype(np.int32) + xmin = np.clip(pts[:, 0].min(), 0, w - 1) + xmax = np.clip(pts[:, 0].max(), 0, w - 1) + ymin = np.clip(pts[:, 1].min(), 0, h - 1) + ymax = np.clip(pts[:, 1].max(), 0, h - 1) + 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 def bitmap_to_boxes( self, diff --git a/doctr/models/preprocessor/pytorch.py b/doctr/models/preprocessor/pytorch.py index ef0d8cf888..d6ef716753 100644 --- a/doctr/models/preprocessor/pytorch.py +++ b/doctr/models/preprocessor/pytorch.py @@ -154,7 +154,6 @@ def __call__( samples = list(multithread_exec(self.sample_transforms, x)) # Batching if self.resize.return_padding_mask: - print(samples) img_batches, mask_batches = self.batch_inputs(samples) else: img_batches = self.batch_inputs(samples)