From 39c45bc9177312dc35dcd673f4b03e27f0fb5e56 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Thu, 24 Jul 2025 08:48:10 +0200
Subject: [PATCH 01/37] docs, type hints,

---
 deeptrack/math.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 13af56ce..b847642b 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1085,8 +1085,8 @@ class Pool(Feature):
     pooling_function: function
         A function that is applied to each local region of the image.
         DOES NOT NEED TO BE WRAPPED IN ANOTHER FUNCTION.
-        The `pooling_function` must accept the input image as a keyword argument
-        named `input`, as it is called via `utils.safe_call`.
+        The `pooling_function` must accept the input image as a keyword
+        argument named `input`, as it is called via `utils.safe_call`.
         Examples include `np.mean`, `np.max`, `np.min`, etc.
     ksize: int
         Size of the pooling kernel.
@@ -1095,7 +1095,8 @@ class Pool(Feature):
 
     Methods
     -------
-    `get(image: np.ndarray | Image, ksize: int, **kwargs: Any) --> np.ndarray`
+    `get(image: NDArray | torch.Tensor | Image,
+         ksize: int, **kwargs: Any) --> NDArray | torch.Tensor`
         Applies the pooling function to the input image.
 
     Examples
@@ -1152,7 +1153,7 @@ def __init__(
 
     def get(
         self: Pool,
-        image: np.ndarray | Image,
+        image: NDArray | torch.Tensor | Image,
         ksize: int,
         **kwargs: Any,
     ) -> np.ndarray:
@@ -1162,7 +1163,7 @@ def get(
 
         Parameters
         ----------
-        image: np.ndarray
+        image: np.ndarray | torch.Tensor | Image
             The input image to pool.
         ksize: int
             Size of the pooling kernel.
@@ -1171,7 +1172,7 @@ def get(
 
         Returns
         -------
-        np.ndarray
+        np.ndarray | torch.Tensor
             The pooled image.
 
         """

From e928b8777d700485fc1c919dd6452ca80509588a Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Thu, 24 Jul 2025 09:18:52 +0200
Subject: [PATCH 02/37] types

---
 deeptrack/math.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index b847642b..4ce4f348 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1156,10 +1156,10 @@ def get(
         image: NDArray | torch.Tensor | Image,
         ksize: int,
         **kwargs: Any,
-    ) -> np.ndarray:
+    ) -> np.ndarray | torch.Tensor:
         """Applies the pooling function to the input image.
 
-        This method applies the pooling function to the input image.
+        This method applies `pooling_function` to the input image.
 
         Parameters
         ----------

From 44f3f72d5b621db1e4e681ab1c4b14ec9620e081 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 10:17:06 +0200
Subject: [PATCH 03/37] docs

---
 deeptrack/math.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 4ce4f348..926150f0 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1156,14 +1156,14 @@ def get(
         image: NDArray | torch.Tensor | Image,
         ksize: int,
         **kwargs: Any,
-    ) -> np.ndarray | torch.Tensor:
+    ) -> NDArray | torch.Tensor:
         """Applies the pooling function to the input image.
 
         This method applies `pooling_function` to the input image.
 
         Parameters
         ----------
-        image: np.ndarray | torch.Tensor | Image
+        image: NDArray | torch.Tensor | Image
             The input image to pool.
         ksize: int
             Size of the pooling kernel.
@@ -1172,7 +1172,7 @@ def get(
 
         Returns
         -------
-        np.ndarray | torch.Tensor
+        NDArray | torch.Tensor
             The pooled image.
 
         """

From 9385fb379467c0535a465094c50cbab38f8a6149 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 10:52:24 +0200
Subject: [PATCH 04/37] torch compatible

---
 deeptrack/math.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 926150f0..d4343ef9 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1176,11 +1176,19 @@ def get(
             The pooled image.
 
         """
+        image_is_torch = isinstance(image, torch.Tensor)
+
+        # Convert to numpy to use skimage.measure.block_reduce.
+        if image_is_torch:
+            device = image.device
+            dtype = image.dtype
+            image = image.detach().cpu().numpy()
 
         kwargs.pop("func", False)
         kwargs.pop("image", False)
         kwargs.pop("block_size", False)
-        return utils.safe_call(
+
+        pooled_image = utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
             func=self.pooling,
@@ -1188,6 +1196,16 @@ def get(
             **kwargs,
         )
 
+        # Convert back to torch.Tensor if needed.
+        if image_is_torch:
+            return torch.tensor(
+                pooled_image,
+                dtype=dtype,
+                device=device
+            )
+        
+        return pooled_image
+
 
 #TODO ***AL*** revise AveragePooling - torch, typing, docstring, unit test
 class AveragePooling(Pool):

From 5123938d05bb387909879abb66d90858401294ac Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:02:33 +0200
Subject: [PATCH 05/37] added averagepooling for torch

---
 deeptrack/math.py | 67 ++++++++++++++++++++++-------------------------
 1 file changed, 31 insertions(+), 36 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index d4343ef9..86b79997 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1078,11 +1078,11 @@ class Pool(Feature):
     non-overlapping blocks of size `ksize` and applying the specified pooling
     function to each block. The result is a downsampled image where each pixel
     value represents the result of the pooling function applied to the
-    corresponding block.
+    corresponding block. This pooling only works with numpy functions.  
 
     Parameters
     ----------
-    pooling_function: function
+    pooling_function:  Numpy function
         A function that is applied to each local region of the image.
         DOES NOT NEED TO BE WRAPPED IN ANOTHER FUNCTION.
         The `pooling_function` must accept the input image as a keyword
@@ -1095,8 +1095,8 @@ class Pool(Feature):
 
     Methods
     -------
-    `get(image: NDArray | torch.Tensor | Image,
-         ksize: int, **kwargs: Any) --> NDArray | torch.Tensor`
+    `get(image: NDArray,
+         ksize: int, **kwargs: Any) --> NDArray`
         Applies the pooling function to the input image.
 
     Examples
@@ -1153,17 +1153,17 @@ def __init__(
 
     def get(
         self: Pool,
-        image: NDArray | torch.Tensor | Image,
+        image: NDArray,
         ksize: int,
         **kwargs: Any,
-    ) -> NDArray | torch.Tensor:
+    ) -> NDArray:
         """Applies the pooling function to the input image.
 
         This method applies `pooling_function` to the input image.
 
         Parameters
         ----------
-        image: NDArray | torch.Tensor | Image
+        image: NDArray | torch.Tensor
             The input image to pool.
         ksize: int
             Size of the pooling kernel.
@@ -1176,19 +1176,11 @@ def get(
             The pooled image.
 
         """
-        image_is_torch = isinstance(image, torch.Tensor)
-
-        # Convert to numpy to use skimage.measure.block_reduce.
-        if image_is_torch:
-            device = image.device
-            dtype = image.dtype
-            image = image.detach().cpu().numpy()
 
         kwargs.pop("func", False)
         kwargs.pop("image", False)
         kwargs.pop("block_size", False)
-
-        pooled_image = utils.safe_call(
+        return utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
             func=self.pooling,
@@ -1196,26 +1188,18 @@ def get(
             **kwargs,
         )
 
-        # Convert back to torch.Tensor if needed.
-        if image_is_torch:
-            return torch.tensor(
-                pooled_image,
-                dtype=dtype,
-                device=device
-            )
-        
-        return pooled_image
-
 
 #TODO ***AL*** revise AveragePooling - torch, typing, docstring, unit test
 class AveragePooling(Pool):
+    # Check if numpy, call super, else use torch.AveragePooling2D
     """Apply average pooling to an image.
 
-    This class reduces the resolution of an image by dividing it into
-    non-overlapping blocks of size `ksize` and applying the average function to
-    each block. The result is a downsampled image where each pixel value
-    represents the average value within the corresponding block of the
-    original image.
+    This class inherits from `Pool` to reduce the resolution of an image by
+    dividing it into non-overlapping blocks of size `ksize` and applying the
+    average function to each block. The result is a downsampled image where
+    each pixel value represents the average value within the corresponding
+    block of the original image. If TORCH_AVAILABLE, it will return the output
+    of `torch.nn.functional.avg_pool2d` instead.
 
     Parameters
     ----------
@@ -1240,10 +1224,12 @@ class AveragePooling(Pool):
 
     Notes
     -----
-    Calling this feature returns a `np.ndarray` by default. If
-    `store_properties` is set to `True`, the returned array will be
-    automatically wrapped in an `Image` object. This behavior is handled
-    internally and does not affect the return type of the `get()` method.
+    Calling this feature returns a pooled image of the input, it will return
+    either numpy or torch depending on the `TORCH_AVAILABLE` flag. If
+    `store_properties` is set to `True` and the input is a numpy array,
+    the returned array will be automatically wrapped in an `Image` object.
+    This behavior is handled internally and does not affect the return type
+    of the `get()` method.
 
     """
 
@@ -1254,7 +1240,8 @@ def __init__(
     ):
         """Initialize the parameters for average pooling.
 
-        This constructor initializes the parameters for average pooling.
+        This constructor initializes the parameters for average pooling and
+        checks whether to use the numpy or torch implementation.
 
         Parameters
         ----------
@@ -1267,6 +1254,14 @@ def __init__(
 
         super().__init__(np.mean, ksize=ksize, **kwargs)
 
+        def get(self, image, ksize: int, **kwargs):
+            
+            # Check torch backend and if the type name starts with "torch".
+            # Isinstance will not work when torch is not available.
+            if TORCH_AVAILABLE and type(image).__module__.startswith("torch"):
+                return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
+            return super().get(image, ksize=ksize, **kwargs)
+
 
 #TODO ***AL*** revise MaxPooling - torch, typing, docstring, unit test
 class MaxPooling(Pool):

From 8d1763ff8c2642f54e1cc5461cfadb44c06985d9 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:03:33 +0200
Subject: [PATCH 06/37] u

---
 deeptrack/math.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 86b79997..80e7cc8b 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1069,7 +1069,6 @@ def __init__(
         super().__init__(ndimage.median_filter, size=ksize, **kwargs)
 
 
-#TODO ***AL*** revise Pool - torch, typing, docstring, unit test
 class Pool(Feature):
     """Downsamples the image by applying a function to local regions of the
     image.
@@ -1189,7 +1188,6 @@ def get(
         )
 
 
-#TODO ***AL*** revise AveragePooling - torch, typing, docstring, unit test
 class AveragePooling(Pool):
     # Check if numpy, call super, else use torch.AveragePooling2D
     """Apply average pooling to an image.
@@ -1255,7 +1253,7 @@ def __init__(
         super().__init__(np.mean, ksize=ksize, **kwargs)
 
         def get(self, image, ksize: int, **kwargs):
-            
+
             # Check torch backend and if the type name starts with "torch".
             # Isinstance will not work when torch is not available.
             if TORCH_AVAILABLE and type(image).__module__.startswith("torch"):

From 30e4c3cba24f87a81df99b1a077ac2c2cdc25cc7 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:16:51 +0200
Subject: [PATCH 07/37] unit test for avgpooling

---
 deeptrack/tests/test_math.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 197afb40..85f1c674 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -109,6 +109,13 @@ def test_AveragePooling(self):
         pooled_image = feature.resolve(input_image)
         self.assertTrue(np.all(pooled_image == [[3.5, 5.5]]))
 
+        if TORCH_AVAILABLE:
+            input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
+                                            5.0, 6.0, 7.0, 8.0]]]])
+            feature = math.AveragePooling(ksize=2)
+            pooled_image = feature.resolve(input_tensor)
+            self.AssertTrue(torch.all(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
+
     def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
         feature = math.MaxPooling(ksize=2)

From e6b94e26ee8d0a00fbb862ff2bd11c21cd8e1966 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:20:33 +0200
Subject: [PATCH 08/37] parenthesis

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 85f1c674..aedbb775 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -111,7 +111,7 @@ def test_AveragePooling(self):
 
         if TORCH_AVAILABLE:
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
-                                            5.0, 6.0, 7.0, 8.0]]]])
+                                          [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
             pooled_image = feature.resolve(input_tensor)
             self.AssertTrue(torch.all(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))

From a88966db076711c419456318149b69e9860530de Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:22:51 +0200
Subject: [PATCH 09/37] input image

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index aedbb775..dff301dc 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature.resolve(input_tensor)
+            pooled_image = feature.resolve(input_image)
             self.AssertTrue(torch.all(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):

From ecf5efcc60432cc90c5483967a525a496c5250ad Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:29:32 +0200
Subject: [PATCH 10/37] capital

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index dff301dc..4dd307e6 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -114,7 +114,7 @@ def test_AveragePooling(self):
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
             pooled_image = feature.resolve(input_image)
-            self.AssertTrue(torch.all(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
+            self.assertTrue(torch.all(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

From 937dc393f731f82d0903a9f1bb353a82b0478e62 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:33:41 +0200
Subject: [PATCH 11/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 4dd307e6..fc436523 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -114,7 +114,7 @@ def test_AveragePooling(self):
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
             pooled_image = feature.resolve(input_image)
-            self.assertTrue(torch.all(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
+            self.assertTrue(torch.allclose(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

From eac7e7ab32d49fcdcca79fbb4abab41c140c70c1 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:38:12 +0200
Subject: [PATCH 12/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index fc436523..5fe1ca63 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -114,7 +114,7 @@ def test_AveragePooling(self):
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
             pooled_image = feature.resolve(input_image)
-            self.assertTrue(torch.allclose(pooled_image == torch.tensor([[[[3.5, 5.5]]]])))
+            self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

From d71e476ebcfb274b4bdda960e6d04f34ac5b53be Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:05:28 +0200
Subject: [PATCH 13/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 5fe1ca63..527e37ff 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature.resolve(input_image)
+            pooled_image = feature(input_image)
             self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):

From 47711838f107e46cacf7298225c9f4d7feb1d34a Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:05:45 +0200
Subject: [PATCH 14/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 527e37ff..0fa33795 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature(input_image)
+            pooled_image = feature.get(input_image)
             self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):

From fbd36c9f1b9ad383f489b3200d90f79394cead30 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:09:56 +0200
Subject: [PATCH 15/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 0fa33795..527e37ff 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature.get(input_image)
+            pooled_image = feature(input_image)
             self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):

From bb107e0918448894c35b1f471de043156724125e Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:12:48 +0200
Subject: [PATCH 16/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 527e37ff..2e815ef7 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature(input_image)
+            pooled_image = feature.get(input_image, ksize=2)
             self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):

From 958d8b411f79db45417e830e49a924206b42f508 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:18:42 +0200
Subject: [PATCH 17/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 2e815ef7..99b9ab93 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature.get(input_image, ksize=2)
+            pooled_image = torch.tensor(feature.get(input_image, ksize=2))
             self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
 
     def test_MaxPooling(self):

From 10097d95d90882f2f788b967981ab83425a7bbdd Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:21:52 +0200
Subject: [PATCH 18/37] u

---
 deeptrack/tests/test_math.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 99b9ab93..b71dd8c8 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -114,7 +114,8 @@ def test_AveragePooling(self):
                                           [5.0, 6.0, 7.0, 8.0]]]])
             feature = math.AveragePooling(ksize=2)
             pooled_image = torch.tensor(feature.get(input_image, ksize=2))
-            self.assertTrue(torch.allclose(pooled_image, torch.tensor([[[[3.5, 5.5]]]])))
+            expected = torch.tensor([[[[3.5, 5.5]]]], dtype=pooled_image.dtype, device=pooled_image.device)
+            self.assertTrue(torch.allclose(pooled_image, expected))
 
     def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

From 2f030c3287961fd8a4c8f9f4fd7b45e17b146805 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:25:58 +0200
Subject: [PATCH 19/37] u

---
 deeptrack/tests/test_math.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index b71dd8c8..276370ff 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -110,11 +110,13 @@ def test_AveragePooling(self):
         self.assertTrue(np.all(pooled_image == [[3.5, 5.5]]))
 
         if TORCH_AVAILABLE:
-            input_image = torch.tensor([[[[1.0, 2.0, 3.0, 4.0],
-                                          [5.0, 6.0, 7.0, 8.0]]]])
+            input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
+                                           [5.0, 6.0, 7.0, 8.0] ]]])
             feature = math.AveragePooling(ksize=2)
             pooled_image = torch.tensor(feature.get(input_image, ksize=2))
+            
             expected = torch.tensor([[[[3.5, 5.5]]]], dtype=pooled_image.dtype, device=pooled_image.device)
+            self.assertEqual(pooled_image.shape, expected.shape)
             self.assertTrue(torch.allclose(pooled_image, expected))
 
     def test_MaxPooling(self):

From 1d872c0dc354eaa80beb1298acf7620a98163bb8 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:29:42 +0200
Subject: [PATCH 20/37] u

---
 deeptrack/tests/test_math.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 276370ff..dd9e5cce 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,10 +113,11 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
                                            [5.0, 6.0, 7.0, 8.0] ]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = torch.tensor(feature.get(input_image, ksize=2))
+            pooled_image = feature.get(input_image, ksize=2)
             
             expected = torch.tensor([[[[3.5, 5.5]]]], dtype=pooled_image.dtype, device=pooled_image.device)
             self.assertEqual(pooled_image.shape, expected.shape)
+
             self.assertTrue(torch.allclose(pooled_image, expected))
 
     def test_MaxPooling(self):

From 3e0cee9e5fef8e3a9cfffb5b63bfe1e17f8f282f Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:33:51 +0200
Subject: [PATCH 21/37] u

---
 deeptrack/math.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 80e7cc8b..86555137 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1256,8 +1256,9 @@ def get(self, image, ksize: int, **kwargs):
 
             # Check torch backend and if the type name starts with "torch".
             # Isinstance will not work when torch is not available.
-            if TORCH_AVAILABLE and type(image).__module__.startswith("torch"):
-                return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
+            if TORCH_AVAILABLE:
+                if type(image).__module__.startswith("torch"):
+                    return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
             return super().get(image, ksize=ksize, **kwargs)
 
 

From c8f6fc0f480ac57f4411efdeb2b0d8723931bc8c Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:33:56 +0200
Subject: [PATCH 22/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index dd9e5cce..1cc8f004 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -115,7 +115,7 @@ def test_AveragePooling(self):
             feature = math.AveragePooling(ksize=2)
             pooled_image = feature.get(input_image, ksize=2)
             
-            expected = torch.tensor([[[[3.5, 5.5]]]], dtype=pooled_image.dtype, device=pooled_image.device)
+            expected = torch.tensor([[[[3.5, 5.5]]]])
             self.assertEqual(pooled_image.shape, expected.shape)
 
             self.assertTrue(torch.allclose(pooled_image, expected))

From 981a7cd4f40f41994acc38fc2f13a9828296e18d Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:49:42 +0200
Subject: [PATCH 23/37] u

---
 deeptrack/math.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 86555137..05c57a46 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1252,14 +1252,13 @@ def __init__(
 
         super().__init__(np.mean, ksize=ksize, **kwargs)
 
-        def get(self, image, ksize: int, **kwargs):
-
-            # Check torch backend and if the type name starts with "torch".
-            # Isinstance will not work when torch is not available.
-            if TORCH_AVAILABLE:
-                if type(image).__module__.startswith("torch"):
-                    return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
-            return super().get(image, ksize=ksize, **kwargs)
+    def get(self, image, ksize: int, **kwargs):
+        # Check torch backend and if the type name starts with "torch".
+        # Isinstance will not work when torch is not available.
+        if TORCH_AVAILABLE:
+            if type(image).__module__.startswith("torch"):
+                return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
+        return super().get(image, ksize=ksize, **kwargs)
 
 
 #TODO ***AL*** revise MaxPooling - torch, typing, docstring, unit test

From 46fb8739ff0acdab2a1eae23b30b4e4aae05fe7e Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 28 Jul 2025 14:53:15 +0200
Subject: [PATCH 24/37] u

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 1cc8f004..dfec1c9e 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -113,7 +113,7 @@ def test_AveragePooling(self):
             input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
                                            [5.0, 6.0, 7.0, 8.0] ]]])
             feature = math.AveragePooling(ksize=2)
-            pooled_image = feature.get(input_image, ksize=2)
+            pooled_image = feature(input_image, ksize=2)
             
             expected = torch.tensor([[[[3.5, 5.5]]]])
             self.assertEqual(pooled_image.shape, expected.shape)

From 3c11ccddb8bad737c4d70b25e95332e3d3f8180c Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Tue, 29 Jul 2025 11:25:53 +0200
Subject: [PATCH 25/37] Update math.py

---
 deeptrack/math.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 05c57a46..97de7eab 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1189,7 +1189,6 @@ def get(
 
 
 class AveragePooling(Pool):
-    # Check if numpy, call super, else use torch.AveragePooling2D
     """Apply average pooling to an image.
 
     This class inherits from `Pool` to reduce the resolution of an image by
@@ -1252,7 +1251,12 @@ def __init__(
 
         super().__init__(np.mean, ksize=ksize, **kwargs)
 
-    def get(self, image, ksize: int, **kwargs):
+    def get(
+        self,
+        image,
+        ksize: int,
+        **kwargs
+    ):
         # Check torch backend and if the type name starts with "torch".
         # Isinstance will not work when torch is not available.
         if TORCH_AVAILABLE:

From a411bbfbdd900d8a15f1b42608b7ea47eafc6829 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Tue, 29 Jul 2025 18:37:20 +0200
Subject: [PATCH 26/37] avg pooling with torch and numpy independent

---
 deeptrack/math.py | 52 ++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 10 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 97de7eab..b5631dd2 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1069,7 +1069,7 @@ def __init__(
         super().__init__(ndimage.median_filter, size=ksize, **kwargs)
 
 
-class Pool(Feature):
+class Pool(Feature): # Deprecated, children will be independent in the future.
     """Downsamples the image by applying a function to local regions of the
     image.
 
@@ -1251,18 +1251,50 @@ def __init__(
 
         super().__init__(np.mean, ksize=ksize, **kwargs)
 
-    def get(
+    def _get_numpy(
         self,
-        image,
-        ksize: int,
+        image: NDArray,
+        ksize: int=3,
+    ):
+        """Method to perform average pooling with the numpy backend enabled.
+
+        Returns the result of the image passed to the scikit image block_reduce
+        function with `np.mean()` as the pooling function.
+        
+        """
+        return utils.safe_call(
+            skimage.measure.block_reduce,
+            image=image,
+            func=self.pooling, # This will be np.mean for this class.
+            block_size=ksize,
+            **kwargs,
+        )
+
+    def _get_torch(
+        self,
+        image: torch.Tensor,
+        ksize: int=3,
+    ):
+        """Method to perform average pooling with the torch backend enabled.
+        
+        Returns the result of the image passed to a torch average pooling layer.   
+
+        """
+
+        return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
+
+     def get(
+        self,
+        image: torch.Tensor | NDArray,
+        ksize: int=3,
         **kwargs
     ):
-        # Check torch backend and if the type name starts with "torch".
-        # Isinstance will not work when torch is not available.
-        if TORCH_AVAILABLE:
-            if type(image).__module__.startswith("torch"):
-                return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
-        return super().get(image, ksize=ksize, **kwargs)
+        if self.backend == "numpy":
+            return self._get_numpy(image, ksize, **kwargs)
+        elif self.backend == "torch":
+            return self._get_torch(image, ksize, **kwargs)
+        else:
+            raise NotImplementedError(f"Backend {self.backend} not supported")
 
 
 #TODO ***AL*** revise MaxPooling - torch, typing, docstring, unit test

From 6f27463f6c27d743bd1ec76a7cff61129ff2d270 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Tue, 29 Jul 2025 20:25:32 +0200
Subject: [PATCH 27/37] docs for averagepooling, unit tests todo

---
 deeptrack/math.py | 68 ++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 59 insertions(+), 9 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index b5631dd2..a7ef54ef 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1254,13 +1254,26 @@ def __init__(
     def _get_numpy(
         self,
         image: NDArray,
-        ksize: int=3,
+        ksize: int = 3,
+        **kwargs,
     ):
         """Method to perform average pooling with the numpy backend enabled.
 
         Returns the result of the image passed to the scikit image block_reduce
         function with `np.mean()` as the pooling function.
-        
+
+        Parameters
+        ----------
+        image: NDArray
+            Input image to be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        NDArray
+            The pooled image as a `NDArray`.
+            
         """
         return utils.safe_call(
             skimage.measure.block_reduce,
@@ -1274,25 +1287,62 @@ def _get_torch(
         self,
         image: torch.Tensor,
         ksize: int=3,
+        **kwargs,
     ):
         """Method to perform average pooling with the torch backend enabled.
         
-        Returns the result of the image passed to a torch average pooling layer.   
+        Returns the result of the image passed to a torch average
+        pooling layer.
+
+        Parameters
+        ----------
+        image: torch.Tensor
+            Input image to be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        torch.Tensor
+            The pooled image as a `torch.Tensor`.
 
         """
 
-        return torch.nn.functional.avg_pool2d(image, kernel_size=ksize)
+        return torch.nn.functional.avg_pool2d(
+            image,
+            kernel_size=ksize, 
+            **kwargs,
+        )
 
-     def get(
+    def get(
         self,
-        image: torch.Tensor | NDArray,
+        image: NDArray | torch.Tensor,
         ksize: int=3,
-        **kwargs
+        **kwargs,
     ):
+        """Method to perform pooling with either torch or numpy backend.
+        
+        Checks the current backend and chooses the appropriate function to pool
+        the input image, either `_get_torch` or `_get_numpy`.
+
+        Parameters
+        ----------
+        image: NDArray | torch.Tensor
+            Input image to be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        NDArray | torch.Tensor
+            The pooled image as `NDArray` or `torch.Tensor` depending on
+            the backend.
+
+        """
         if self.backend == "numpy":
-            return self._get_numpy(image, ksize, **kwargs)
+            return self._get_numpy(image, ksize, **kwargs,)
         elif self.backend == "torch":
-            return self._get_torch(image, ksize, **kwargs)
+            return self._get_torch(image, ksize, **kwargs,)
         else:
             raise NotImplementedError(f"Backend {self.backend} not supported")
 

From 00fa7957c4b48c8066653d02c7966c3fda1a938b Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 09:58:44 +0200
Subject: [PATCH 28/37] self.get_backend

---
 deeptrack/math.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index a7ef54ef..31655c56 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1339,9 +1339,9 @@ def get(
             the backend.
 
         """
-        if self.backend == "numpy":
+        if self.get_backend() == "numpy":
             return self._get_numpy(image, ksize, **kwargs,)
-        elif self.backend == "torch":
+        elif self.get_backend() == "torch":
             return self._get_torch(image, ksize, **kwargs,)
         else:
             raise NotImplementedError(f"Backend {self.backend} not supported")

From a1e839c004fa08dfdf7023b2b7069e0bee42db6a Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 10:03:11 +0200
Subject: [PATCH 29/37] separated numpy and torch unit tests

---
 deeptrack/tests/test_math.py | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index dfec1c9e..fe55d198 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -82,12 +82,27 @@ def test_Blur(self):
         #blurred_image = feature.resolve(input_image)
         #self.assertTrue(xp.all(blurred_image == expected_output))
 
+    def test_AveragePooling(self):
+        input_image = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=float)
+        feature = math.AveragePooling(ksize=2)
+        pooled_image = feature.resolve(input_image)
+        self.assertTrue(np.all(pooled_image == [[3.5, 5.5]]))
+
 
 # Extending the test and setting the backend to torch
 @unittest.skipUnless(TORCH_AVAILABLE, "PyTorch is not installed.")
 class TestMath_Torch(TestMath_Numpy):
     BACKEND = "torch"
-    pass
+
+    input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
+                                    [5.0, 6.0, 7.0, 8.0] ]]])
+    feature = math.AveragePooling(ksize=2)
+    pooled_image = feature(input_image, ksize=2)
+    expected = torch.tensor([[[[3.5, 5.5]]]])
+    self.assertEqual(pooled_image.shape, expected.shape)
+    self.assertTrue(torch.allclose(pooled_image, expected))
+
+    
 
 
 class TestMath(unittest.TestCase):
@@ -109,16 +124,6 @@ def test_AveragePooling(self):
         pooled_image = feature.resolve(input_image)
         self.assertTrue(np.all(pooled_image == [[3.5, 5.5]]))
 
-        if TORCH_AVAILABLE:
-            input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
-                                           [5.0, 6.0, 7.0, 8.0] ]]])
-            feature = math.AveragePooling(ksize=2)
-            pooled_image = feature(input_image, ksize=2)
-            
-            expected = torch.tensor([[[[3.5, 5.5]]]])
-            self.assertEqual(pooled_image.shape, expected.shape)
-
-            self.assertTrue(torch.allclose(pooled_image, expected))
 
     def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

From fe6b1bc0955f8733ad45dd43121ab994c9cde5dc Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 10:04:32 +0200
Subject: [PATCH 30/37] u

---
 deeptrack/tests/test_math.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index fe55d198..9baf3962 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -93,14 +93,15 @@ def test_AveragePooling(self):
 @unittest.skipUnless(TORCH_AVAILABLE, "PyTorch is not installed.")
 class TestMath_Torch(TestMath_Numpy):
     BACKEND = "torch"
-
-    input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
-                                    [5.0, 6.0, 7.0, 8.0] ]]])
-    feature = math.AveragePooling(ksize=2)
-    pooled_image = feature(input_image, ksize=2)
-    expected = torch.tensor([[[[3.5, 5.5]]]])
-    self.assertEqual(pooled_image.shape, expected.shape)
-    self.assertTrue(torch.allclose(pooled_image, expected))
+    
+    def test_AveragePooling(self):
+        input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
+                                        [5.0, 6.0, 7.0, 8.0] ]]])
+        feature = math.AveragePooling(ksize=2)
+        pooled_image = feature(input_image, ksize=2)
+        expected = torch.tensor([[[[3.5, 5.5]]]])
+        self.assertEqual(pooled_image.shape, expected.shape)
+        self.assertTrue(torch.allclose(pooled_image, expected))
 
     
 

From 5b94e5a0666d58f85ae7e6ed23234366b562fd95 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 10:16:03 +0200
Subject: [PATCH 31/37] attempt at fixing broken optics

---
 deeptrack/optics.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/deeptrack/optics.py b/deeptrack/optics.py
index 803504f2..fb81ff5b 100644
--- a/deeptrack/optics.py
+++ b/deeptrack/optics.py
@@ -155,6 +155,8 @@ def _pad_volume(
 from deeptrack.image import Image, pad_image_to_fft, maybe_cupy
 from deeptrack.types import ArrayLike, PropertyLike
 
+import torch
+
 from deeptrack import image
 from deeptrack import units_registry as u
 
@@ -358,6 +360,8 @@ def get(
             )
 
             imaged_sample = self._objective.resolve(sample_volume)
+            if self.get_backend() == "torch": 
+                imaged_sample = torch.from_numpy(imaged_sample)
 
         # Upscale given by the optics needs to be handled separately.
         if _upscale_given_by_optics != (1, 1, 1):

From 2332197bd6c87d66c47164d884ac61a1eefca755 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 10:19:42 +0200
Subject: [PATCH 32/37] u

---
 deeptrack/optics.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/deeptrack/optics.py b/deeptrack/optics.py
index fb81ff5b..267a7e00 100644
--- a/deeptrack/optics.py
+++ b/deeptrack/optics.py
@@ -155,7 +155,10 @@ def _pad_volume(
 from deeptrack.image import Image, pad_image_to_fft, maybe_cupy
 from deeptrack.types import ArrayLike, PropertyLike
 
-import torch
+from deeptrack.backend import config, TORCH_AVAILABLE, xp
+
+if TORCH_AVAILABLE:
+    import torch
 
 from deeptrack import image
 from deeptrack import units_registry as u

From 3c611592d3ff489f795dc7f98be231e631f10e9d Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 10:22:54 +0200
Subject: [PATCH 33/37] undo fix

---
 deeptrack/optics.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/deeptrack/optics.py b/deeptrack/optics.py
index 267a7e00..803504f2 100644
--- a/deeptrack/optics.py
+++ b/deeptrack/optics.py
@@ -155,11 +155,6 @@ def _pad_volume(
 from deeptrack.image import Image, pad_image_to_fft, maybe_cupy
 from deeptrack.types import ArrayLike, PropertyLike
 
-from deeptrack.backend import config, TORCH_AVAILABLE, xp
-
-if TORCH_AVAILABLE:
-    import torch
-
 from deeptrack import image
 from deeptrack import units_registry as u
 
@@ -363,8 +358,6 @@ def get(
             )
 
             imaged_sample = self._objective.resolve(sample_volume)
-            if self.get_backend() == "torch": 
-                imaged_sample = torch.from_numpy(imaged_sample)
 
         # Upscale given by the optics needs to be handled separately.
         if _upscale_given_by_optics != (1, 1, 1):

From a24973764810bd890a172ffda9807c5be358d766 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 15:51:51 +0200
Subject: [PATCH 34/37] no kwargs in torch pool

---
 deeptrack/math.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 31655c56..49f9e4ad 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1310,8 +1310,7 @@ def _get_torch(
 
         return torch.nn.functional.avg_pool2d(
             image,
-            kernel_size=ksize, 
-            **kwargs,
+            kernel_size=ksize,
         )
 
     def get(

From c8d2fb86b71c28a0197ae1d4be0a15413fd557ce Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 16:01:27 +0200
Subject: [PATCH 35/37] docs typo

---
 deeptrack/math.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 49f9e4ad..d6cd624d 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1195,8 +1195,8 @@ class AveragePooling(Pool):
     dividing it into non-overlapping blocks of size `ksize` and applying the
     average function to each block. The result is a downsampled image where
     each pixel value represents the average value within the corresponding
-    block of the original image. If TORCH_AVAILABLE, it will return the output
-    of `torch.nn.functional.avg_pool2d` instead.
+    block of the original image. If the backend is torch, it will return the
+    output of `torch.nn.functional.avg_pool2d` instead.
 
     Parameters
     ----------
@@ -1222,7 +1222,7 @@ class AveragePooling(Pool):
     Notes
     -----
     Calling this feature returns a pooled image of the input, it will return
-    either numpy or torch depending on the `TORCH_AVAILABLE` flag. If
+    either numpy or torch depending on the backend. If
     `store_properties` is set to `True` and the input is a numpy array,
     the returned array will be automatically wrapped in an `Image` object.
     This behavior is handled internally and does not affect the return type

From 920712624042059e1ab1615ffabf077ccd364acb Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 16:02:39 +0200
Subject: [PATCH 36/37] docs

---
 deeptrack/math.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index d6cd624d..a778559d 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1238,7 +1238,8 @@ def __init__(
         """Initialize the parameters for average pooling.
 
         This constructor initializes the parameters for average pooling and
-        checks whether to use the numpy or torch implementation.
+        checks whether to use the numpy or torch implementation, defaults to
+        numpy.
 
         Parameters
         ----------

From 288d944fbd4cb8dd5a17a25d6b8640330d2cffb9 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 6 Aug 2025 15:24:35 -0700
Subject: [PATCH 37/37] Added shape handling for len(dim) == 2

---
 deeptrack/math.py | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index a778559d..ff81fe6e 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1193,10 +1193,15 @@ class AveragePooling(Pool):
 
     This class inherits from `Pool` to reduce the resolution of an image by
     dividing it into non-overlapping blocks of size `ksize` and applying the
-    average function to each block. The result is a downsampled image where
-    each pixel value represents the average value within the corresponding
-    block of the original image. If the backend is torch, it will return the
-    output of `torch.nn.functional.avg_pool2d` instead.
+    `max` function to each block. The result is a downsampled image where each
+    pixel value represents the maximum value within the corresponding block of
+    the original image. This is useful for reducing the size of an image while
+    retaining the most significant features.
+
+    If the backend is numpy, the downsampling is performed using
+    `skimage.measure.block_reduce`.
+    If the backend is torch, the downsampling
+    is performed using `torch.nn.functional.avg_pool2d`.
 
     Parameters
     ----------
@@ -1222,11 +1227,10 @@ class AveragePooling(Pool):
     Notes
     -----
     Calling this feature returns a pooled image of the input, it will return
-    either numpy or torch depending on the backend. If
-    `store_properties` is set to `True` and the input is a numpy array,
-    the returned array will be automatically wrapped in an `Image` object.
-    This behavior is handled internally and does not affect the return type
-    of the `get()` method.
+    either numpy or torch depending on the backend. If `store_properties` is
+    set to `True` and the input is a numpy array, the returned array will be
+    automatically wrapped in an `Image` object. This behavior is handled
+    internally and does not affect the return type of the `get()` method.
 
     """
 
@@ -1309,6 +1313,16 @@ def _get_torch(
 
         """
 
+        # If needed, expand tensor shape
+        if len(image.shape) == 2:
+            expanded_image = image.unsqueeze(0)
+
+            pooled_image = torch.nn.functional.avg_pool2d(
+                expanded_image, kernel_size=ksize,
+            )
+            # Remove the expanded dim.
+            return pooled_image.squeeze(0)
+
         return torch.nn.functional.avg_pool2d(
             image,
             kernel_size=ksize,