from __future__ import annotations
from collections.abc import Sequence
from typing import Literal, Self, cast
import cv2
import numpy as np
from neuralib.typing import PathLike, PathLikeType
__all__ = ['image_array',
'ImageArrayWrapper']
RGB_CHANNEL_TYPE = Literal['r', 'g', 'b', 'red', 'green', 'blue']
[docs]
def image_array(dat: np.ndarray | PathLike, *,
mode: Literal['RGB', 'RGBA', 'gray'] | None = None,
alpha: bool = False) -> ImageArrayWrapper:
"""
Image array numpy subclass
:param dat: Image data as a NumPy array or a file path
:param mode: Color mode {'RGB', 'RGBA', 'gray'}. Optional if dat is ndarray.
:param alpha: If True and loading from file, convert to RGBA
:return:
"""
return ImageArrayWrapper(dat, mode=mode, alpha=alpha)
[docs]
class ImageArrayWrapper(np.ndarray):
"""Subclass of numpy.ndarray that wraps an image and provides chainable image processing methods"""
[docs]
def __new__(cls, dat: np.ndarray | PathLike, *,
mode: Literal['RGB', 'RGBA', 'gray'] | None = None,
alpha: bool = False) -> Self:
"""
:param dat: Image data as a NumPy array or a file path
:param mode: Color mode {'RGB', 'RGBA', 'gray'}. Optional if dat is ndarray.
:param alpha: If True and loading from file, convert to RGBA
"""
if isinstance(dat, PathLikeType):
img = cv2.imread(str(dat))
if img is None:
raise FileNotFoundError(str(dat))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if alpha:
img = cv2.cvtColor(img, cv2.COLOR_RGB2RGBA)
mode = 'RGBA'
else:
mode = 'RGB'
arr = cast(np.ndarray, img)
else:
arr = np.asarray(dat)
match (mode, arr.ndim, arr.shape[-1] if arr.ndim == 3 else None):
case (None, 2, _):
mode = 'gray'
case (None, 3, 3):
mode = 'RGB'
case (None, 3, 4):
mode = 'RGBA'
case ('gray', 2, _):
pass # valid grayscale
case ('RGB', 3, 3):
pass # valid rgb
case ('RGBA', 3, 4):
pass # valid rgba
case _:
raise ValueError(f"Invalid shape {arr.shape} for mode={mode!r}")
obj = arr.view(cls)
obj.mode = mode
return cast(Self, obj)
def __array_finalize__(self, obj):
if obj is None:
return
self.mode = getattr(obj, 'mode', 'RGB')
@property
def height(self) -> int:
"""image height"""
return self.shape[0]
@property
def width(self) -> int:
"""image width"""
return self.shape[1]
[docs]
def to_gray(self) -> ImageArrayWrapper:
"""convert the image array to grayscale"""
if self.ndim == 2:
ret = self.copy().astype("uint8")
elif self.shape[2] == 3:
ret = cv2.cvtColor(self, cv2.COLOR_RGB2GRAY)
elif self.shape[2] == 4:
ret = cv2.cvtColor(self, cv2.COLOR_RGBA2GRAY)
else:
raise RuntimeError(f'Unexpected image shape: {self.shape} for grayscale conversion')
return ImageArrayWrapper(ret, mode='gray')
[docs]
def flipud(self) -> ImageArrayWrapper:
"""flip the image array upside down (vertical)"""
return ImageArrayWrapper(np.flipud(self))
[docs]
def fliplr(self) -> ImageArrayWrapper:
"""flip the image array left to right (horizontal flip)"""
return ImageArrayWrapper(np.fliplr(self))
[docs]
def select_channel(self, channel: RGB_CHANNEL_TYPE) -> ImageArrayWrapper:
"""extract a single color channel from an RGB or RGBA image
:param channel: one of 'r'/'red', 'g'/'green', or 'b'/'blue'.
"""
if self.ndim < 3:
raise RuntimeError(f'Image shape invalid for splitting channel: {self.ndim}')
channels = cv2.split(self)
match channel:
case 'r' | 'red':
ret = channels[0]
case 'g' | 'green':
ret = channels[1]
case 'b' | 'blue':
ret = channels[2]
case _:
raise ValueError(f'invalid {channel} argument')
return ImageArrayWrapper(ret, mode='gray')
[docs]
def view_2d(self, flipud: bool = False) -> ImageArrayWrapper:
"""
Convert a multi-channel image to a 2D representation.
- For a 4-channel image, the array is reinterpreted as a 2D array of 32-bit integers.
- For a 3-channel image, the result is a grayscale image obtained by applying luminance conversion.
:param flipud: reverse the order of elements along axis 0 (up/down)
"""
if self.ndim == 2:
return self
match self.shape[2]:
case 4:
w, h, _ = self.shape
ret = self.view(dtype=np.uint32).reshape((w, h))
case 3:
r, g, b = self[:, :, 0], self[:, :, 1], self[:, :, 2]
ret = (0.2989 * r + 0.5870 * g + 0.1140 * b).astype(np.uint8)
case _:
raise ValueError(f'invalid arr shape: {self.shape}')
zelf = ImageArrayWrapper(ret, mode='gray')
if flipud:
zelf = zelf.flipud()
return zelf
[docs]
def gaussian_blur(self, ksize: Sequence[int], sigma_x: float, sigma_y: float, **kwargs) -> ImageArrayWrapper:
"""
Apply a Gaussian blur to the image.
:param ksize: Kernel size (e.g., (5, 5)). The width and height should be odd numbers.
:param sigma_x: Standard deviation in the X direction.
:param sigma_y: Standard deviation in the Y direction.
:param kwargs: Additional keyword arguments for ``cv2.GaussianBlur()``.
"""
img = cv2.GaussianBlur(self, ksize=ksize, sigmaX=sigma_x, sigmaY=sigma_y, **kwargs)
return ImageArrayWrapper(img)
[docs]
def canny_filter(self, threshold_1: float = 30, threshold_2: float = 150, **kwargs) -> ImageArrayWrapper:
"""
Apply the Canny edge detection algorithm to the grayscale version of the image.
:param threshold_1: The first threshold for the hysteresis procedure.
:param threshold_2: The second threshold for the hysteresis procedure.
:param kwargs: Additional keyword arguments for ``cv2.Canny()``.
"""
img = cv2.Canny(self.to_gray(), threshold1=threshold_1, threshold2=threshold_2, **kwargs)
return ImageArrayWrapper(img)
[docs]
def binarize(self, thresh: float, maxval: float = 255, **kwargs) -> ImageArrayWrapper:
"""
Convert the image to a binary image using a fixed threshold.
:param thresh: Threshold value. Pixels above this value are set to maxval; otherwise, 0.
:param maxval: The value to use for pixels above the threshold.
:param kwargs: Additional keyword arguments for ``cv2.threshold()``.
"""
_, img = cv2.threshold(self, thresh, maxval=maxval, type=cv2.THRESH_BINARY, **kwargs)
return ImageArrayWrapper(img)
[docs]
def denoise(self, h: int = 10, temp_win_size: int = 7, search_win_size: int = 21, **kwargs) -> ImageArrayWrapper:
"""
Apply Non-local Means Denoising to the image.
- For grayscale images, cv2.fastNlMeansDenoising is used.
- For color images, cv2.fastNlMeansDenoisingColored is used.
:param h: Filtering parameter controlling the degree of smoothing.
:param temp_win_size: Template window size in pixels.
:param search_win_size: Search window size in pixels.
:param kwargs: Additional keyword arguments for the cv2 denoising function.
"""
if self.mode == 'gray':
fn = cv2.fastNlMeansDenoising
else:
fn = cv2.fastNlMeansDenoisingColored
img = fn(self, h=h, templateWindowSize=temp_win_size, searchWindowSize=search_win_size, **kwargs)
return ImageArrayWrapper(img)
[docs]
def enhance_contrast(self) -> ImageArrayWrapper:
"""Enhance the contrast of the image using histogram equalization"""
match self.mode:
case 'gray':
eq = cv2.equalizeHist(self)
return ImageArrayWrapper(eq, mode='gray')
case 'RGB':
ycrcb = cv2.cvtColor(self, cv2.COLOR_RGB2YCrCb)
ycrcb[:, :, 0] = cv2.equalizeHist(ycrcb[:, :, 0])
eq = cv2.cvtColor(ycrcb, cv2.COLOR_YCrCb2RGB)
return ImageArrayWrapper(eq, mode='RGB')
case 'RGBA':
rgb = self[..., :3]
alpha = self[..., 3]
ycrcb = cv2.cvtColor(rgb, cv2.COLOR_RGB2YCrCb)
ycrcb[:, :, 0] = cv2.equalizeHist(ycrcb[:, :, 0])
eq_rgb = cv2.cvtColor(ycrcb, cv2.COLOR_YCrCb2RGB)
eq = np.dstack((eq_rgb, alpha))
return ImageArrayWrapper(eq, mode='RGBA')
case _:
raise ValueError(f'invalid mode: {self.mode}')
[docs]
def local_maxima(self, channel: RGB_CHANNEL_TYPE, **kwargs) -> ImageArrayWrapper:
"""
Compute the local maxima of the image on a specified color channel.
The specified channel is first extracted (and returned as a grayscale image),
then the skimage local_maxima function is applied.
:param channel: one of 'r'/'red', 'g'/'green', or 'b'/'blue'.
:param kwargs: additional keyword arguments for ``skimage.morphology.local_maxima()``.
"""
from skimage.morphology import local_maxima
img = self.select_channel(channel)
if np.sum(img) == 0:
return ImageArrayWrapper(np.zeros_like(img, dtype=np.uint8), mode='gray')
else:
return ImageArrayWrapper(np.asarray(local_maxima(img, **kwargs)), mode='gray')