From ab373486fd4e0f31b11831dbb0c5ad76e1124da6 Mon Sep 17 00:00:00 2001
From: Elliot Young <elliot.young1996@gmail.com>
Date: Wed, 29 May 2024 00:44:09 -0700
Subject: [PATCH] Fix methods in `image_util.py`

---
 src/frheed/image_util.py | 51 +++++++++++++++++++++++++++-------------
 1 file changed, 35 insertions(+), 16 deletions(-)

diff --git a/src/frheed/image_util.py b/src/frheed/image_util.py
index f59ad56..065b99c 100644
--- a/src/frheed/image_util.py
+++ b/src/frheed/image_util.py
@@ -8,6 +8,24 @@
 ImageArray = npt.NDArray[np.generic]
 
 
+def get_image_shape(image: ImageArray) -> tuple[int, int, int]:
+    """Returns the height, width, and number of channels of an image.
+
+    Raises:
+        TypeError if the image has less than 2 dimensions or greater than 3 dimensions.
+    """
+    match image.ndim:
+        case 2:
+            height, width = image.shape
+            channels = 1
+        case 3:
+            height, width, channels = image.shape
+        case _:
+            raise TypeError(f"Image with shape {image.shape} has unsupported number of dimensions")
+
+    return (height, width, channels)
+
+
 def qimage_to_ndarray(image: QtGui.QImage) -> ImageArray:
     """Converts a QImage to a numpy array."""
     if (image_bits := image.bits()) is None:
@@ -24,7 +42,9 @@ def qimage_to_ndarray(image: QtGui.QImage) -> ImageArray:
 
 def get_rectangle_region(image: ImageArray, x1: int, y1: int, x2: int, y2: int) -> ImageArray:
     """Returns the region of an image bound by a rectangle."""
-    return np.copy(image[y1:y2, x1:x2])
+    # The region must have nonzero dimensions
+    h, w = image.shape[:2]
+    return np.copy(image[min(y1, h - 2) : max(y2, y1 + 1), min(x1, w - 2) : max(x2, x1 + 1)])
 
 
 def get_ellipse_region(image: ImageArray, x1: int, y1: int, x2: int, y2: int) -> ImageArray:
@@ -32,11 +52,11 @@ def get_ellipse_region(image: ImageArray, x1: int, y1: int, x2: int, y2: int) ->
     # Get the region within the ellipse bounding box
     region = get_rectangle_region(image, x1, y1, x2, y2)
 
-    # Create an elliptical mask where 1 = pixel to include, 0 = pixel to exclude
-    height, width = region.shape[:2]
+    # Mask all pixels outside the ellipse
+    height, width, channels = get_image_shape(region)
     center_x = width // 2
     center_y = height // 2
-    mask = np.zeros((height, width), dtype=np.uint8)
+    mask = np.ones(region.shape, dtype=np.uint8)
     cv2.ellipse(
         mask,
         center=(center_x, center_y),
@@ -44,12 +64,10 @@ def get_ellipse_region(image: ImageArray, x1: int, y1: int, x2: int, y2: int) ->
         angle=0,
         startAngle=0,
         endAngle=360,
-        color=(1,),
+        color=[0] * channels,
         thickness=-1,  # negative thickness will draw a filled ellipse
     )
-
-    # Set all pixels outside the elliptical mask to 0
-    return cv2.bitwise_and(region, region, mask=mask)
+    return np.ma.MaskedArray(region, mask=np.ma.make_mask(mask))
 
 
 def get_line_region(
@@ -58,11 +76,12 @@ def get_line_region(
     """Returns the region of an image under a line."""
     # Get the region within the line bounding box
     region = get_rectangle_region(image, x1, y1, x2, y2)
-
-    # Create a mask by drawing the line
-    height, width = region.shape[:2]
-    mask = np.zeros((height, width), dtype=np.uint8)
-    cv2.line(region, pt1=(0, 0), pt2=(height, width), color=(1,), thickness=thickness)
-
-    # Set all pixels not under the line to 0
-    return cv2.bitwise_and(region, region, mask=mask)
+    if x1 == x2 or y1 == y2:
+        # Region is already a line and does not require masking
+        return region
+
+    # Mask all pixels not under the line
+    height, width, channels = get_image_shape(region)
+    mask = np.zeros(region.shape, dtype=np.uint8)
+    cv2.line(mask, pt1=(0, 0), pt2=(height, width), color=[0] * channels, thickness=thickness)
+    return np.ma.MaskedArray(region, mask=np.ma.make_mask(mask))