Mg/features plot

deeptrack/features.py

@@ -387,7 +387,19 @@ class Feature(DeepTrackNode):
         behavior.
     `bind_arguments(arguments: Feature) -> Feature`
         It binds another feature’s properties as arguments to this feature.
-    `plot(input_image: np.ndarray | list[np.ndarray] | Image | list[Image] | None = None, resolve_kwargs: dict | None = None, interval: float | None = None, **kwargs: Any) -> Any`
+    `plot(
+        input_image: (
+            NDArray
+            | list[NDArray]
+            | torch.Tensor
+            | list[torch.Tensor]
+            | Image
+            | list[Image]
+        ) = None,
+        resolve_kwargs: dict | None = None,
+        interval: float | None = None,
+        **kwargs: Any,
+    ) -> Any`
         It visualizes the output of the feature.
 
     **Private and internal methods.**
@@ -1754,54 +1766,91 @@ def bind_arguments(
 
         return self
 
-    #TODO ***MG***
     def plot(
         self: Feature,
-        input_image: np.ndarray | list[np.ndarray] | Image | list[Image] = None,
+        input_image: (
+            NDArray
+            | list[NDArray]
+            | torch.Tensor
+            | list[torch.Tensor]
+            | Image
+            | list[Image]
+        ) = None,
         resolve_kwargs: dict = None,
         interval: float = None,
         **kwargs: Any,
     ) -> Any:
-        """Visualizes the output of the feature.
+        """Visualize the output of the feature.
 
-        This method resolves the feature and visualizes the result. If the output is 
-        an `Image`, it displays it using `pyplot.imshow`. If the output is a list, it 
-        creates an animation. In Jupyter notebooks, the animation is played inline 
-        using `to_jshtml()`. In scripts, the animation is displayed using the 
+        `plot()` resolves the feature and visualizes the result. If the output
+        is a single image (NumPy array, PyTorch tensor, or Image), it is
+        displayed using `pyplot.imshow`. If the output is a list, an animation
+        is created. In Jupyter notebooks, the animation is played inline using
+        `to_jshtml()`. In scripts, the animation is displayed using the
         matplotlib backend.
 
         Any parameters in `kwargs` are passed to `pyplot.imshow`.
 
         Parameters
         ----------
-        input_image: np.ndarray or Image or list[np.ndarray or Image], optional
-            The input image or list of images passed as an argument to the `resolve` 
-            call. If `None`, uses previously set input values or propagates properties.
+        input_image: np.ndarray, torch.tensor, or Image or list[np.ndarray,
+            torch.tensor, or Image], optional
+            The input image or list of images passed as an argument to the
+            `resolve` call. If `None`, uses previously set input values or
+            propagates properties.
         resolve_kwargs: dict, optional
             Additional keyword arguments passed to the `resolve` call.
         interval: float, optional
-            The time between frames in the animation, in milliseconds. The default 
-            value is 33 ms.
+            The time between frames in the animation, in milliseconds. The
+            default value is 33 ms.
         **kwargs: dict, optional
             Additional keyword arguments passed to `pyplot.imshow`.
-       
+
         Returns
         -------
         Any
             The output of the feature or pipeline after execution.
+
+        Examples
+        --------
+        >>> import deeptrack as dt
+
+        Create an instance of a dummy feature that returns the input:
+        >>> feature = dt.DummyFeature()
+
+        Generate and plot a grayscale image:
+        >>> import numpy as np
+        >>>
+        >>> img = np.random.randint(0, 256, (64, 64))
+        >>> feature.plot(img, cmap="gray");
+
+        Generate and plot a grayscale video:
+        >>> video = [np.random.randint(0, 256, (64, 64)) for _ in range(10)]
+        >>> feature.plot(video, interval=100, cmap="gray");
+
+        Generate a grayscale image using torch and plot it:
+        >>> import torch
+        >>>
+        >>> img = torch.randint(0, 256, size=(64, 64))
+        >>> feature.plot(img, cmap="gray");
+
+        Generate a simulated image of a point particle visualized using
+        brightfield microscopy and plot it:
+        >>> particle = dt.PointParticle()
+        >>> optics = dt.Brightfield()
+        >>> imaged_particle = optics(particle)
+        >>> imaged_particle.plot(cmap="gray");
 
         """
 
         from IPython.display import HTML, display
 
-        # if input_image is not None:
-        #     input_image = [Image(input_image)]
-
         output_image = self.resolve(input_image, **(resolve_kwargs or {}))
 
         # If a list, assume video
         if not isinstance(output_image, list):
             # Single image
+            output_image = xp.squeeze(output_image)
             plt.imshow(output_image, **kwargs)
             return plt.gca()
 
@@ -1810,11 +1859,14 @@ def plot(
         images = []
         plt.axis("off")
         for image in output_image:
+            image = xp.squeeze(image)
             images.append([plt.imshow(image, **kwargs)])
 
         if not interval:
             if isinstance(output_image[0], Image):
-                interval = output_image[0].get_property("interval") or (1 / 30 * 1000)
+                interval = (
+                    output_image[0].get_property("interval") or (1 / 30 * 1000)
+                )
             else:
                 interval = 1 / 30 * 1000