fix: Fixed Resize transfo in PyTorch when aspect ratio matches the ta…

…rget (#357) * fix: Fixed edge case of torch resize * test: Added unittest to ensure correct behaviour
mindee · Jul 6, 2021 · d7f50b5 · d7f50b5
1 parent 0445126
commit d7f50b5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/doctr/transforms/modules/pytorch.py b/doctr/transforms/modules/pytorch.py
@@ -22,32 +22,34 @@ def __init__(
         preserve_aspect_ratio: bool = False,
         symmetric_pad: bool = False,
     ) -> None:
-        super().__init__(size[::-1], interpolation)
+        super().__init__(size, interpolation)
         self.preserve_aspect_ratio = preserve_aspect_ratio
         self.symmetric_pad = symmetric_pad
 
     def forward(self, img: torch.Tensor) -> torch.Tensor:
-        target_ratio = self.size[1] / self.size[0]
+        target_ratio = self.size[0] / self.size[1]
         actual_ratio = img.shape[-2] / img.shape[-1]
         if not self.preserve_aspect_ratio or (target_ratio == actual_ratio):
             return super().forward(img)
         else:
+            # Resize
             if actual_ratio > target_ratio:
-                tmp_size = (self.size[1], int(self.size[1] / actual_ratio))
+                tmp_size = (self.size[0], int(self.size[0] / actual_ratio))
             else:
-                tmp_size = (int(self.size[0] * actual_ratio), self.size[0])
+                tmp_size = (int(self.size[1] * actual_ratio), self.size[1])
+
             # Scale image
             img = F.resize(img, tmp_size, self.interpolation)
             # Pad (inverted in pytorch)
-            _pad = (0, self.size[0] - img.shape[-1], 0, self.size[1] - img.shape[-2])
+            _pad = (0, self.size[1] - img.shape[-1], 0, self.size[0] - img.shape[-2])
             if self.symmetric_pad:
                 half_pad = (math.ceil(_pad[1] / 2), math.ceil(_pad[3] / 2))
                 _pad = (half_pad[0], _pad[1] - half_pad[0], half_pad[1], _pad[3] - half_pad[1])
             return pad(img, _pad)
 
     def __repr__(self) -> str:
         interpolate_str = self.interpolation.value
-        _repr = f"output_size={self.size[::-1]}, interpolation='{interpolate_str}'"
+        _repr = f"output_size={self.size}, interpolation='{interpolate_str}'"
         if self.preserve_aspect_ratio:
             _repr += f", preserve_aspect_ratio={self.preserve_aspect_ratio}, symmetric_pad={self.symmetric_pad}"
         return f"{self.__class__.__name__}({_repr})"
diff --git a/test/pytorch/test_transforms_pt.py b/test/pytorch/test_transforms_pt.py
@@ -40,6 +40,12 @@ def test_resize():
     assert not torch.all(out == 1)
     assert out.shape[-2:] == output_size
 
+    # Same aspect ratio
+    output_size = (32, 128)
+    transfo = Resize(output_size, preserve_aspect_ratio=True)
+    out = transfo(torch.ones((3, 16, 64), dtype=torch.float32))
+    assert out.shape[-2:] == output_size
+
 
 @pytest.mark.parametrize(
     "rgb_min",