Update of noises/background

deeptrack/noises.py

@@ -63,21 +63,59 @@ class Noise(Feature):
     """Base abstract noise class."""
 
 
-#TODO ***MG*** revise Background - torch, typing, docstring, unit test
 class Background(Noise):
-    """Adds a constant value to an image
+    """Adds a constant value to an image.
 
     Parameters
     ----------
     offset : float
-        The value to add to the image
+        The value to add to the image.
+    **kwargs : Any
+        Additional keyword arguments passed to the parent `Noise` class.
+
+    Methods
+    -------
+    get(
+        image: np.ndarray, torch.Tensor, or Image,
+        offset: float,
+        **kwargs,
+        ) -> np.ndarray, torch.Tensor, or Image
+        Adds the constant offset to the input image.
+
+    Examples
+    --------
+    >>> import deeptrack as dt
+    >>> import numpy as np
+
+    Create an input image with zeros:
+    >>> input_image = np.zeros((2,2))
+
+    Define the Background noise feature with offset 0.5:
+    >>> noise = dt.Background(offset=0.5)
+
+    Apply the noise to the input image and print the resulting image:
+    >>> output_image = noise.resolve(input_image)
+    >>> print(output_image)
+    [[0.5 0.5]
+    [0.5 0.5]]
+
     """
 
     def __init__(
         self: Background,
         offset: PropertyLike[float],
         **kwargs: Any,
     ):
+        """
+        Initialize the Background noise feature.
+
+        Parameters
+        ----------
+        offset : PropertyLike[float]
+            The constant value to be added to the image.
+        **kwargs : Any
+            Additional arguments passed to the parent `Noise` class.
+        """
         super().__init__(offset=offset, **kwargs)
 
     def get(
@@ -86,6 +124,21 @@ def get(
         offset: float,
         **kwargs: Any,
     ) -> NDArray[Any] | torch.Tensor | Image:
+        """
+        Add the given offset to the image.
+
+        Parameters
+        ----------
+        image : np.ndarray, torch.Tensor, or Image
+            The input image.
+        offset : float
+            The value to add to the image.
+
+        Returns
+        -------
+        np.ndarray, torch.Tensor, or Image
+            The image with offset added.
+        """
 
         return image + offset
 

deeptrack/tests/test_noises.py

@@ -7,26 +7,49 @@
 import numpy as np
 
 from deeptrack.image import Image
-from deeptrack import noises
+from deeptrack import noises, TORCH_AVAILABLE
 
+if TORCH_AVAILABLE:
+    import torch
 
 class TestNoises(unittest.TestCase):
     def test_Offset(self):
         noise = noises.Offset(offset=0.5)
         input_image = Image(np.zeros((256, 256)))
         output_image = noise.resolve(input_image)
+
         self.assertIsInstance(output_image, np.ndarray)
         self.assertEqual(output_image.shape, (256, 256))
         self.assertTrue(np.all(np.array(output_image) == 0.5))
 
     def test_Background(self):
+        # Test with DeepTrack Image
         noise = noises.Background(offset=0.5)
         input_image = Image(np.zeros((256, 256)))
         output_image = noise.resolve(input_image)
+
         self.assertIsInstance(output_image, np.ndarray)
         self.assertEqual(output_image.shape, (256, 256))
         self.assertTrue(np.all(np.array(output_image) == 0.5))
 
+        # Test with NumPy array
+        noise = noises.Background(offset=0.5)
+        input_image = np.ones((10, 10))
+        output_image = noise.resolve(input_image)
+        self.assertIsInstance(output_image, np.ndarray)
+        self.assertEqual(output_image.shape, (10, 10))
+        self.assertTrue(np.all(np.array(output_image) == 1.5))
+
+        ### Test with PyTorch tensor (if available)
+        if TORCH_AVAILABLE:
+            noise = noises.Background(offset=0.25)
+            input_image = torch.zeros(5,5)
+            output_image = noise.resolve(input_image)
+
+            self.assertIsInstance(output_image, torch.Tensor)
+            self.assertEqual(output_image.shape, (5,5))
+            self.assertTrue(torch.all(output_image == 0.25).item())
+
     def test_Gaussian(self):
         noise = noises.Gaussian(mu=0.1, sigma=0.05)
         input_image = Image(np.zeros((256, 256)))