diff --git a/deeptrack/features.py b/deeptrack/features.py index a4bf2c709..5443f8496 100644 --- a/deeptrack/features.py +++ b/deeptrack/features.py @@ -216,7 +216,7 @@ def propagate_data_to_dependencies( "Merge", "OneOf", "OneOfDict", - "LoadImage", # TODO ***MG*** + "LoadImage", "SampleToMasks", # TODO ***MG*** "AsType", # TODO ***MG*** "ChannelFirst2d", @@ -7195,21 +7195,21 @@ def get( class LoadImage(Feature): """Load an image from disk and preprocess it. - This feature loads an image file using multiple fallback file readers - (`imageio`, `numpy`, `Pillow`, and `OpenCV`) until a suitable reader is - found. The image can be optionally converted to grayscale, reshaped to - ensure a minimum number of dimensions, or treated as a list of images if + `LoadImage` loads an image file using multiple fallback file readers + (`imageio`, `numpy`, `Pillow`, and `OpenCV`) until a suitable reader is + found. The image can be optionally converted to grayscale, reshaped to + ensure a minimum number of dimensions, or treated as a list of images if multiple paths are provided. Parameters ---------- path: PropertyLike[str or list[str]] - The path(s) to the image(s) to load. Can be a single string or a list + The path(s) to the image(s) to load. Can be a single string or a list of strings. load_options: PropertyLike[dict[str, Any]], optional Additional options passed to the file reader. It defaults to `None`. as_list: PropertyLike[bool], optional - If `True`, the first dimension of the image will be treated as a list. + If `True`, the first dimension of the image will be treated as a list. It defaults to `False`. ndim: PropertyLike[int], optional Ensures the image has at least this many dimensions. It defaults to @@ -7217,7 +7217,7 @@ class LoadImage(Feature): to_grayscale: PropertyLike[bool], optional If `True`, converts the image to grayscale. It defaults to `False`. get_one_random: PropertyLike[bool], optional - If `True`, extracts a single random image from a stack of images. Only + If `True`, extracts a single random image from a stack of images. Only used when `as_list` is `True`. It defaults to `False`. Attributes @@ -7228,7 +7228,15 @@ class LoadImage(Feature): Methods ------- - `get(image: Any, path: str or list[str], load_options: dict[str, Any] | None, ndim: int, to_grayscale: bool, as_list: bool, get_one_random: bool, **kwargs: Any) -> array` + `get( + path: str | list[str], + load_options: dict[str, Any] | None, + ndim: int, + to_grayscale: bool, + as_list: bool, + get_one_random: bool, + **kwargs: Any, + ) -> NDArray | list[NDArray] | torch.Tensor | list[torch.Tensor]` Load the image(s) from disk and process them. Raises @@ -7236,6 +7244,12 @@ class LoadImage(Feature): IOError If no file reader could parse the file or the file does not exist. + Notes + ---- + By default, `LoadImage` returns a NumPy array. If you want the output as + a PyTorch tensor, convert the feature to torch by calling `.torch()` before + resolving. + Examples -------- >>> import deeptrack as dt @@ -7243,7 +7257,7 @@ class LoadImage(Feature): Create a temporary image file: >>> import numpy as np >>> import os, tempfile - >>> + >>> >>> temp_file = tempfile.NamedTemporaryFile(suffix=".npy", delete=False) >>> np.save(temp_file.name, np.random.rand(100, 100, 3)) @@ -7271,7 +7285,14 @@ class LoadImage(Feature): ... ) >>> loaded_image = load_image_feature.resolve() >>> loaded_image.shape - (2, 2, 3, 1) + (100, 100, 3, 1) + + Load an image as a PyTorch tensor by setting the backend of the feature: + >>> load_image_feature = dt.LoadImage(path=temp_file.name) + >>> load_image_feature.torch() + >>> loaded_image = load_image_feature.resolve() + >>> type(loaded_image) + Cleanup the temporary file: >>> os.remove(temp_file.name) @@ -7313,7 +7334,7 @@ def __init__( If `True`, selects a single random image from a stack when `as_list=True`. It defaults to `False`. **kwargs: Any - Additional keyword arguments passed to the parent `Feature` class, + Additional keyword arguments passed to the parent `Feature` class, allowing further customization. """ @@ -7338,31 +7359,36 @@ def get( as_list: bool, get_one_random: bool, **kwargs: Any, - ) -> NDArray | torch.Tensor: + ) -> NDArray[Any] | torch.Tensor | list: """Load and process an image or a list of images from disk. - This method attempts to load an image using multiple file readers - (`imageio`, `numpy`, `Pillow`, and `OpenCV`) until a valid format is + This method attempts to load an image using multiple file readers + (`imageio`, `numpy`, `Pillow`, and `OpenCV`) until a valid format is found. It supports optional processing steps such as ensuring a minimum - number of dimensions, grayscale conversion, and treating multi-frame + number of dimensions, grayscale conversion, and treating multi-frame images as lists. + The output is returned as a NumPy array by default. If `as_list=True`, + the result is a Python list of arrays. If the backend of the feature is + `"torch"`, the image is returned as a PyTorch tensor. + Parameters ---------- path: str or list[str] - The file path(s) to the image(s) to be loaded. A single string + The file path(s) to the image(s) to be loaded. A single string loads one image, while a list of paths loads multiple images. load_options: dict of str to Any, optional - Additional options passed to the file reader (e.g., `allow_pickle` + Additional options passed to the file reader (e.g., `allow_pickle` for NumPy, `mode` for OpenCV). It defaults to `None`. ndim: int - Ensures the image has at least this many dimensions. If the loaded - image has fewer dimensions, extra dimensions are added. + Ensures the image has at least this many dimensions. If the loaded + image has fewer dimensions, extra dimensions are added. It defaults + to `3`. to_grayscale: bool If `True`, converts the image to grayscale. It defaults to `False`. as_list: bool - If `True`, treats the first dimension as a list of images instead - of stacking them into a NumPy array. + If `True`, treats the first dimension as a list of images instead + of stacking them into a NumPy array. It defaults to `False`. get_one_random: bool If `True`, selects a single random image from a multi-frame stack when `as_list=True`. It defaults to `False`. @@ -7372,14 +7398,14 @@ def get( Returns ------- array - The loaded and processed image(s). If `as_list=True`, returns a + The loaded and processed image(s). If `as_list=True`, returns a list of images; otherwise, returns a single NumPy array or PyTorch tensor. Raises ------ IOError - If no valid file reader is found or if the specified file does not + If no valid file reader is found or if the specified file does not exist. """ @@ -7402,8 +7428,9 @@ def get( try: import PIL.Image - image = [PIL.Image.open(file, **load_options) - for file in path] + image = [ + PIL.Image.open(file, **load_options) for file in path + ] except (IOError, ImportError): import cv2 @@ -7439,11 +7466,18 @@ def get( ) # Ensure the image has at least `ndim` dimensions. - while ndim and image.ndim < ndim: - image = np.expand_dims(image, axis=-1) + if not isinstance(image, list) and ndim: + while image.ndim < ndim: + image = np.expand_dims(image, axis=-1) # Convert to PyTorch tensor if needed. - #TODO + if self.get_backend() == "torch": + + # Convert to stack if needed. + if isinstance(image, list): + image = np.stack(image, axis=0) + + image = torch.from_numpy(image) return image diff --git a/deeptrack/tests/test_features.py b/deeptrack/tests/test_features.py index 591547a68..6d4a9aa88 100644 --- a/deeptrack/tests/test_features.py +++ b/deeptrack/tests/test_features.py @@ -1836,9 +1836,14 @@ def test_LoadImage(self): try: with NamedTemporaryFile(suffix=".npy", delete=False) as temp_npy: - np.save(temp_npy.name, test_image_array) + pass + np.save(temp_npy.name, test_image_array) # npy_filename = temp_npy.name + with NamedTemporaryFile(suffix=".npy", delete=False) as temp_npy2: + pass + np.save(temp_npy2.name, test_image_array) + with NamedTemporaryFile(suffix=".png", delete=False) as temp_png: PIL_Image.fromarray(test_image_array).save(temp_png.name) # png_filename = temp_png.name @@ -1877,12 +1882,61 @@ def test_LoadImage(self): loaded_image = load_feature.resolve() self.assertGreaterEqual(len(loaded_image.shape), 4) + # Test loading a list of images + load_feature = features.LoadImage( + path=[temp_npy.name, temp_npy2.name], as_list=True + ) + loaded_list = load_feature.resolve() + self.assertIsInstance(loaded_list, list) + self.assertEqual(len(loaded_list), 2) + + for img in loaded_list: + self.assertTrue(isinstance(img, np.ndarray)) + + # Test loading a random image from a list of images + load_feature = features.LoadImage( + path=[temp_npy.name, temp_npy2.name], + ndim=4, + as_list=True, + get_one_random=True, + ) + loaded_image = load_feature.resolve() + self.assertTrue( + np.allclose( + loaded_image[:, :, 0, 0], test_image_array, rtol=1.e-3 + ) + ) + self.assertEqual(loaded_image.shape, (50, 50, 1, 1)) + + import gc + gc.collect() + + # Test loading an image as a torch tensor. + if TORCH_AVAILABLE: + load_feature = features.LoadImage(path=temp_png.name) + load_feature.torch() + loaded_image = load_feature.resolve() + self.assertIsInstance(loaded_image, torch.Tensor) + self.assertEqual( + loaded_image.shape[:2], test_image_array.shape + ) + + loaded_image_np = loaded_image.numpy() + self.assertTrue( + np.allclose( + test_image_array, loaded_image_np[:, :, 0], rtol=1.e-3 + ) + ) + finally: - for file in [temp_npy.name, temp_png.name, temp_jpg.name]: + for file in [ + temp_npy.name, + temp_png.name, + temp_jpg.name, + temp_npy2.name + ]: os.remove(file) - #TODO: Add a test for loading a list of images. - def test_SampleToMasks(self): # Parameters