diff --git a/deeptrack/math.py b/deeptrack/math.py index 13af56cec..41d89abbb 100644 --- a/deeptrack/math.py +++ b/deeptrack/math.py @@ -1444,75 +1444,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: diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py index 197afb404..02ff4235e 100644 --- a/deeptrack/tests/test_math.py +++ b/deeptrack/tests/test_math.py @@ -136,9 +136,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):