Add image data classes

plantcv/data/__init__.py

@@ -0,0 +1,147 @@
+"""PlantCV data classes"""
+import numpy as np
+
+
+class Image(np.ndarray):
+    """Generic image class that extends the np.ndarray class."""
+
+    # From NumPy documentation
+    # Add uri attribute
+    def __new__(cls, input_array: np.ndarray, uri: str):
+        """
+        Create a new Image instance.
+
+        Parameters
+        ----------
+        input_array : np.ndarray
+            The input array representing the image data.
+        uri : str
+            Uniform resource identifier of the source file.
+
+        Returns
+        -------
+        Image
+            An instance of the Image class with the uri attribute set.
+        """
+        obj = np.asarray(input_array).view(cls)
+        # New attribute uri stores uniform resource identifier of the source file
+        obj.uri = uri
+        return obj
+
+    def __array_finalize__(self, obj):
+        if obj is not None:
+            self.uri = getattr(obj, "uri", None)
+
+    def __getitem__(self, key):
+        # Enhance the np.ndarray __getitem__ method
+        # Slice the array as requested but return an array of the same class
+        # Idea from NumPy examples of subclassing:
+        return super(Image, self).__getitem__(key)
+
+
+class GRAY(Image):
+    """Subclass of Image for grayscale images."""
+
+    def __new__(cls, input_array: np.ndarray, uri: str):
+        return Image.__new__(cls, input_array, uri)
+
+
+class BGR(Image):
+    """Subclass of Image for Blue, Green, Red (BGR) images."""
+
+    def __new__(cls, input_array: np.ndarray, uri: str):
+        return Image.__new__(cls, input_array, uri)
+
+    def __getitem__(self, key):
+        # Overwrite the __getitem__ method to return a GRAY object if the
+        # requested slice is 2D
+        new_arr = super(Image, self).__getitem__(key)
+        if len(new_arr.shape) == 2:
+            return GRAY(input_array=new_arr, uri=self.uri)
+        return new_arr
+
+
+class RGB(Image):
+    """Subclass of Image for Red, Green, Blue (RGB) images."""
+
+    def __new__(cls, input_array: np.ndarray, uri: str):
+        return Image.__new__(cls, input_array, uri)
+
+    def __getitem__(self, key):
+        # Overwrite the __getitem__ method to return a GRAY object if the
+        # requested slice is 2D
+        new_arr = super(Image, self).__getitem__(key)
+        if len(new_arr.shape) == 2:
+            return GRAY(input_array=new_arr, uri=self.uri)
+        return new_arr
+
+
+class HSI(Image):
+    """Subclass of Image for hyperspectral images."""
+
+    def __new__(cls, input_array: np.ndarray, uri: str, wavelengths: list, default_wavelengths: list,
+                wavelength_units: str = "nm", metadata: dict = None):
+        # Create an instance of Image with default attributes
+        obj = Image.__new__(cls, input_array, uri)
+        # Add HSI-specific attributes
+        # Set wavelengths list
+        obj.wavelengths = wavelengths
+        # Set wavelength units
+        obj.wavelength_units = wavelength_units
+        # Compute min and max wavelengths
+        obj.min_wavelength = np.min(wavelengths)
+        obj.max_wavelength = np.max(wavelengths)
+        # Set default wavelengths for RGB thumbnail
+        obj.default_wavelengths = default_wavelengths
+        # If no default wavelengths are provided, set to common RGB wavelengths if available,
+        if default_wavelengths is None:
+            if obj.max_wavelength >= 635 and obj.min_wavelength <= 490:
+                obj.default_wavelengths = [480, 540, 710]
+            else:
+                obj.default_wavelengths = [wavelengths[np.argmax(np.sum(input_array, axis=(0, 1)))]]
+        # Set metadata
+        obj.metadata = metadata if metadata is not None else {}
+        return obj
+
+    def get_wavelength(self, wavelength):
+        """
+        Get a specific wavelength from the hyperspectral image.
+
+        Parameters
+        ----------
+        wavelength : float
+            The wavelength to retrieve.
+
+        Returns
+        -------
+        plantcv.data.HSI
+            A new HSI object containing only the specified wavelength.
+        """
+        # Find the index of the closest wavelength
+        idx = np.abs(np.array(self.wavelengths) - wavelength).argmin()
+        # Use the parent class __getitem__ method to get the specific wavelength slice
+        obj = super(HSI, self).__getitem__(np.s_[:, :, idx])
+        # Update attributes for the new object
+        obj.wavelengths = [self.wavelengths[idx]]
+        obj.min_wavelength = np.min(obj.wavelengths)
+        obj.max_wavelength = np.max(obj.wavelengths)
+        return obj
+
+    def thumbnail(self):
+        """
+        Create an RGB or grayscale thumbnail for quick visualization.
+
+        Returns
+        -------
+        plantcv.data.BGR or plantcv.data.GRAY
+            An RGB or grayscale thumbnail representation of the HSI data.
+        """
+        if len(self.default_wavelengths) == 3:
+            thumb = BGR(input_array=np.dstack([self.get_wavelength(self.default_wavelengths[0]),
+                                               self.get_wavelength(self.default_wavelengths[1]),
+                                               self.get_wavelength(self.default_wavelengths[2])]),
+                        uri=self.uri)
+        else:
+            thumb = GRAY(input_array=self.get_wavelength(self.default_wavelengths[0]),
+                         uri=self.uri)
+        return thumb

tests/data/test_data_classes.py

@@ -0,0 +1,87 @@
+import numpy as np
+from plantcv.data import Image, GRAY, BGR, RGB, HSI
+
+
+def test_image():
+    """Test creating an Image class image."""
+    img = Image(input_array=np.zeros((10, 10), dtype=np.uint8), uri="image.png")
+    assert isinstance(img, Image)
+
+
+def test_image_none():
+    """Test creating an Image class image."""
+    img = Image(input_array=None, uri=None)
+    assert isinstance(img, Image)
+
+
+def test_image_slice():
+    """Test subsetting an Image."""
+    img = Image(input_array=np.zeros((10, 10), dtype=np.uint8), uri="image.png")
+    assert img[0:5, 0:5].shape == (5, 5)
+
+
+def test_bgr():
+    """Test creating a BGR class image."""
+    bgr = BGR(input_array=np.zeros((10, 10, 3), dtype=np.uint8), uri="bgr.png")
+    assert isinstance(bgr, BGR)
+
+
+def test_bgr_slice():
+    """Test subsetting a BGR image."""
+    bgr = BGR(input_array=np.zeros((10, 10, 3), dtype=np.uint8), uri="bgr.png")
+    assert bgr[0:5, 0:5].shape == (5, 5, 3)
+
+
+def test_rgb():
+    """Test creating a RGB class image."""
+    bgr = RGB(input_array=np.zeros((10, 10, 3), dtype=np.uint8), uri="rgb.png")
+    assert isinstance(bgr, RGB)
+
+
+def test_rgb_slice():
+    """Test subsetting an RGB image."""
+    rgb = RGB(input_array=np.zeros((10, 10, 3), dtype=np.uint8), uri="rgb.png")
+    assert rgb[0:5, 0:5].shape == (5, 5, 3)
+
+
+def test_gray():
+    """Test creating a GRAY class image."""
+    gray = GRAY(input_array=np.zeros((10, 10), dtype=np.uint8), uri="gray.png")
+    assert isinstance(gray, GRAY)
+
+
+def test_bgr_to_gray():
+    """Test converting a BGR image to a GRAY image."""
+    bgr = BGR(input_array=np.zeros((10, 10, 3), dtype=np.uint8), uri="bgr.png")
+    gray = bgr[:, :, 0]
+    assert isinstance(gray, GRAY)
+
+
+def test_rgb_to_gray():
+    """Test converting a RGB image to a GRAY image."""
+    rgb = RGB(input_array=np.zeros((10, 10, 3), dtype=np.uint8), uri="rgb.png")
+    gray = rgb[:, :, 0]
+    assert isinstance(gray, GRAY)
+
+
+def test_hsi():
+    """Test creating an HSI class image."""
+    hsi = HSI(input_array=np.zeros((10, 10, 5), dtype=np.uint8), uri="hsi.data", wavelengths=[480, 540, 710, 800, 900],
+              default_wavelengths=None, wavelength_units="nm")
+    assert isinstance(hsi, HSI)
+
+
+def test_hsi_grayscale_thumb():
+    """Test creating an HSI class image."""
+    hsi = HSI(input_array=np.zeros((10, 10, 5), dtype=np.uint8), uri="hsi.data", wavelengths=[700, 800, 900],
+              default_wavelengths=None, wavelength_units="nm")
+    thumb = hsi.thumbnail()
+    assert isinstance(thumb, GRAY)
+
+
+def test_hsi_rgb_thumb():
+    """Test creating an HSI class image."""
+    hsi = HSI(input_array=np.zeros((10, 10, 5), dtype=np.uint8), uri="hsi.data", wavelengths=[480, 540, 710, 800, 900],
+              default_wavelengths=None, wavelength_units="nm")
+    thumb = hsi.thumbnail()
+    assert isinstance(thumb, BGR)