diff --git a/plantcv/data/__init__.py b/plantcv/data/__init__.py new file mode 100644 index 000000000..bc2d5c683 --- /dev/null +++ b/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 diff --git a/tests/data/__init__.py b/tests/data/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/test_data_classes.py b/tests/data/test_data_classes.py new file mode 100644 index 000000000..e452c9676 --- /dev/null +++ b/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)