Update math/resize (#406)

* update math/resize

* formatting

* update unittesting

* update resize

* update test resize

* update test resize

* update resize

* update resize

* update resize

Author: mirjagranfors

deeptrack/math.py

@@ -1663,75 +1663,154 @@ def __init__(
         super().__init__(np.median, ksize=ksize, **kwargs)
 
 
-#TODO ***MG*** revise Resize - torch, typing, docstring, unit test
 class Resize(Feature):
     """Resize an image to a specified size.
 
-    This class is a wrapper around cv2.resize and resizes an image to a
-    specified size. The `dsize` parameter specifies the desired output size of
-    the image.
-    Note that the order of the axes is different in cv2 and numpy. In cv2, the
-    first axis is the vertical axis, while in numpy it is the horizontal axis.
-    This is reflected in the default values of the arguments.
+    `Resize` resizes an image using:
+      - OpenCV (`cv2.resize`) for NumPy arrays.
+      - PyTorch (`torch.nn.functional.interpolate`) for PyTorch tensors.
+
+    The interpretation of the `dsize` parameter follows the convention 
+    of the underlying backend:
+      - **NumPy (OpenCV)**: `dsize` is given as `(width, height)` to match
+        OpenCV’s default.
+      - **PyTorch**: `dsize` is given as `(height, width)`.
 
     Parameters
     ----------
-    dsize: tuple
-        Size to resize to.
+    dsize: PropertyLike[tuple[int, int]]
+        The target size. Format depends on backend: `(width, height)` for
+        NumPy, `(height, width)` for PyTorch.
     **kwargs: Any
-        Additional parameters sent to the resizing function.
+        Additional parameters sent to the underlying resize function:
+          - NumPy: passed to `cv2.resize`.
+          - PyTorch: passed to `torch.nn.functional.interpolate`.
+
+    Methods
+    -------
+    get(
+        image: np.ndarray | torch.Tensor, dsize: tuple[int, int], **kwargs
+    ) -> np.ndarray | torch.Tensor
+        Resize the input image to the specified size.
+
+    Examples
+    --------
+    >>> import deeptrack as dt
+
+    Numpy example:
+    >>> import numpy as np
+    >>>
+    >>> input_image = np.random.rand(16, 16)            # Create image
+    >>> feature = dt.math.Resize(dsize=(8, 4))          # (width=8, height=4)
+    >>> resized_image = feature.resolve(input_image)    # Resize it to (4, 8)
+    >>> print(resized_image.shape)
+    (4, 8)
+
+    PyTorch example:
+    >>> import torch
+    >>>
+    >>> input_image = torch.rand(1, 1, 16, 16)          # Create image
+    >>> feature = dt.math.Resize(dsize=(4, 8))          # (height=4, width=8)
+    >>> resized_image = feature.resolve(input_image)    # Resize it to (4, 8)
+    >>> print(resized_image.shape)
+    torch.Size([1, 1, 4, 8])
 
     """
 
     def __init__(
         self: Resize,
-        dsize: PropertyLike[tuple] = (256, 256),
+        dsize: PropertyLike[tuple[int, int]] = (256, 256),
         **kwargs: Any,
     ):
-        """Initialize the parameters for resizing input features.
-
-        This constructor initializes the parameters for resizing input
-        features.
+        """Initialize the parameters for the Resize feature.
 
         Parameters
         ----------
-        dsize: tuple
-            Size to resize to.
+        dsize: PropertyLike[tuple[int, int]]
+            The target size. Format depends on backend: `(width, height)` for
+            NumPy, `(height, width)` for PyTorch. Default is (256, 256).
         **kwargs: Any
-            Additional keyword arguments.
+            Additional arguments passed to the parent `Feature` class.
 
         """
 
         super().__init__(dsize=dsize, **kwargs)
 
-    def get(self: Resize, image: np.ndarray, dsize: tuple, **kwargs: Any) -> np.ndarray:
+    def get(
+        self: Resize,
+        image: NDArray | torch.Tensor,
+        dsize: tuple[int, int],
+        **kwargs: Any,
+    ) -> NDArray | torch.Tensor:
         """Resize the input image to the specified size.
 
-        This method resizes the input image to the specified size.
-
         Parameters
         ----------
-        image: np.ndarray
+        image: np.ndarray or torch.Tensor
             The input image to resize.
-        dsize: tuple
+            - NumPy arrays may be grayscale (H, W) or color (H, W, C).
+            - Torch tensors are expected in one of the following formats:
+              (N, C, H, W), (C, H, W), or (H, W).
+        dsize: tuple[int, int]
             Desired output size of the image.
+            - NumPy: (width, height)
+            - PyTorch: (height, width)
         **kwargs: Any
-            Additional keyword arguments.
+            Additional keyword arguments passed to the underlying resize 
+            function (`cv2.resize` or `torch.nn.functional.interpolate`).
 
         Returns
         -------
-        np.ndarray
-            The resized image.
+        np.ndarray or torch.Tensor
+            The resized image in the same type and dimensionality format as
+            input.
 
-        """
+        Notes
+        -----
+        - For PyTorch tensors, resizing uses bilinear interpolation with
+          `align_corners=False`. This choice matches OpenCV’s `cv2.resize`
+          default behavior when resizing NumPy arrays, aiming to produce nearly
+          identical results between both backends.
 
-        import cv2
-        from deeptrack import config
+        """
 
         if self._wrap_array_with_image:
             image = strip(image)
 
-        return utils.safe_call(cv2.resize, positional_args=[image, dsize], **kwargs)
+        if apc.is_torch_array(image):
+            original_shape = image.shape
+
+            # Reshape input to (N, C, H, W)
+            if image.ndim == 2:     # (H, W)
+                image = image.unsqueeze(0).unsqueeze(0)
+            elif image.ndim == 3:   # (C, H, W)
+                image = image.unsqueeze(0)
+            elif image.ndim != 4:
+                raise ValueError(
+                    "Resize only supports tensors with shape (N, C, H, W), "
+                    "(C, H, W), or (H, W)."
+                )
+
+            resized = torch.nn.functional.interpolate(
+                image,
+                size=dsize,
+                mode="bilinear",
+                align_corners=False,
+            )
+
+            # Restore original dimensionality
+            if len(original_shape) == 2:
+                resized = resized.squeeze(0).squeeze(0)
+            elif len(original_shape) == 3:
+                resized = resized.squeeze(0)
+
+            return resized
+
+        else:
+            import cv2
+            return utils.safe_call(
+                cv2.resize, positional_args=[image, dsize], **kwargs
+            )
 
 
 if OPENCV_AVAILABLE:

deeptrack/tests/test_math.py

@@ -157,9 +157,40 @@ def test_MedianPooling(self):
     @unittest.skipUnless(OPENCV_AVAILABLE, "OpenCV is not installed.")
     def test_Resize(self):
         input_image = np.random.rand(16, 16)
-        feature = math.Resize(dsize=(8, 8))
+        feature = math.Resize(dsize=(8, 4))
         resized = feature.resolve(input_image)
-        self.assertEqual(resized.shape, (8, 8))
+
+        self.assertIsInstance(resized, np.ndarray)
+        self.assertEqual(resized.shape, (4, 8))
+
+    @unittest.skipUnless(TORCH_AVAILABLE, "PyTorch is not installed.")
+    def test_Resize_torch(self):
+
+        feature = math.Resize(dsize=(4, 8))
+
+        input_image = torch.rand(16, 16)
+        resized = feature.resolve(input_image)
+        self.assertIsInstance(resized, torch.Tensor)
+        self.assertEqual(tuple(resized.shape), (4, 8))
+
+        if OPENCV_AVAILABLE:
+            # Compare with NumPy version:
+            feature_np = math.Resize(dsize=(8, 4))
+            input_image_np = input_image.numpy()
+            resized_np = feature_np.resolve(input_image_np)
+            np.testing.assert_allclose(
+                        resized_np, resized.numpy(), rtol=1e-5, atol=1e-5
+                    )
+
+        input_image = torch.rand(3, 16, 16)
+        resized = feature.resolve(input_image)
+        self.assertIsInstance(resized, torch.Tensor)
+        self.assertEqual(tuple(resized.shape), (3, 4, 8))
+
+        input_image = torch.rand(1, 1, 16, 16)
+        resized = feature.resolve(input_image)
+        self.assertIsInstance(resized, torch.Tensor)
+        self.assertEqual(tuple(resized.shape), (1, 1, 4, 8))
 
     @unittest.skipUnless(OPENCV_AVAILABLE, "OpenCV is not installed.")
     def test_BlurCV2_GaussianBlur(self):