FAST tune polling size (mindee#1500)

odulcy-mindee · Mar 5, 2024 · 058b5db · 058b5db
1 parent fa68ac0
commit 058b5db
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/doctr/models/detection/fast/pytorch.py b/doctr/models/detection/fast/pytorch.py
@@ -58,9 +58,9 @@ def __init__(
         out_channels: int = 128,
     ) -> None:
         super().__init__()
-        self.reduction = nn.ModuleList(
-            [FASTConvLayer(in_channels * scale, out_channels, kernel_size=3) for scale in [1, 2, 4, 8]]
-        )
+        self.reduction = nn.ModuleList([
+            FASTConvLayer(in_channels * scale, out_channels, kernel_size=3) for scale in [1, 2, 4, 8]
+        ])
 
     def _upsample(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return F.interpolate(x, size=y.shape[-2:], mode="bilinear")
@@ -122,7 +122,7 @@ def __init__(
         bin_thresh: float = 0.3,
         box_thresh: float = 0.1,
         dropout_prob: float = 0.1,
-        pooling_size: int = 4,  # different from paper performs better on close text-rich images
+        pooling_size: int = 9,
         assume_straight_pages: bool = True,
         exportable: bool = False,
         cfg: Optional[Dict[str, Any]] = {},

diff --git a/doctr/models/detection/fast/tensorflow.py b/doctr/models/detection/fast/tensorflow.py
@@ -125,7 +125,7 @@ def __init__(
         bin_thresh: float = 0.3,
         box_thresh: float = 0.1,
         dropout_prob: float = 0.1,
-        pooling_size: int = 4,  # different from paper performs better on close text-rich images
+        pooling_size: int = 9,
         assume_straight_pages: bool = True,
         exportable: bool = False,
         cfg: Optional[Dict[str, Any]] = {},