Ez lmao

assignment1/solution.py

@@ -386,7 +386,8 @@ def two_e(image: npt.NDArray[np.uint8]) -> None:
 
 def excercise_three() -> None:
     #three_a()
-    three_b()
+    mask1, mask2 = three_b()
+    three_c(uz.imread('./images/bird.jpg'), mask1)
 
 
 def three_a() -> None:
@@ -442,7 +443,7 @@ def three_a() -> None:
     plt.show()
 
 
-def three_b() -> None:
+def three_b():
     """
     Try to clean up the mask of the image bird.jpg using morphological operations
     as shown in the image. Experiment with different sizes of the structuring element.
@@ -457,6 +458,7 @@ def three_b() -> None:
     SE_CROSS = cv2.getStructuringElement(cv2.MORPH_CROSS, (4, 4))
 
     # Perform operation of opening and then closing (almost)
+    # Values found through trial and error
     mask1 = cv2.dilate(mask1, SE_ELIPSE, iterations = 4)
     mask1 = cv2.erode(mask1, SE_ELIPSE)
     mask1 = cv2.dilate(mask1, SE_ELIPSE)
@@ -479,6 +481,24 @@ def three_b() -> None:
     axs[2].imshow(mask2, cmap='gray')
     axs[2].set(title="Using cross")
 
+
+    plt.show()
+
+    return (mask1, mask2)
+
+# Ez lmao
+def three_c(image: npt.NDArray[np.float64], mask: npt.NDArray[np.uint8]):
+    """
+    Write a function immask that accepts a three channel image and
+    a binary mask and returns an image where pixel values are set to black if the
+    corresponding pixel in the mask is equal to 0. Otherwise, the pixel value should be
+    equal to the corresponding image pixel. Do not use for loops, as they are slow
+    """
+    mask = np.expand_dims(mask, axis=2)
+    
+    image = mask * image
+
+    plt.imshow(image)
     plt.show()