diff --git a/assignment1/solution.py b/assignment1/solution.py
index ffd5219..658c493 100644
--- a/assignment1/solution.py
+++ b/assignment1/solution.py
@@ -4,7 +4,6 @@ from matplotlib import pyplot as plt
 import random
 import cv2
 import uz_framework.image as uz_image
-import UZ_utils as uz
 
 #######################################
 # EXCERCISE 1: Basic image processing #
@@ -146,7 +145,7 @@ def excercise_two() -> None:
     two_b('./images/bird.jpg', 100, 20)
     two_c('./images/bird.jpg', 20, 100)
     two_d()
-    two_e(uz.imread_gray('./images/bird.jpg', uz.ImageType.uint8).astype(np.uint8))
+    two_e(uz_image.imread_gray('./images/bird.jpg', uz_image.ImageType.uint8).astype(np.uint8))
 
 
 def two_a() -> tuple[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
@@ -295,9 +294,9 @@ def two_e(image: npt.NDArray[np.uint8]):
 ######################################################
 
 def excercise_three() -> None:
-    three_a()
+    #three_a()
     #mask1, _ = three_b()
-    #three_c(uz.imread('./images/bird.jpg', uz.ImageType.float64), mask1)
+    #three_c(uz_image.imread('./images/bird.jpg', uz_image.ImageType.float64), mask1)
     #three_d()
     three_e()
 
@@ -312,7 +311,7 @@ def three_a() -> None:
     Answer: Opening = Erosion then dialation
             Closing = Dialation then erosion
     """
-    img_orig = uz.imread_gray('./images/mask.png', uz.ImageType.float64)
+    img_orig = uz_image.imread_gray('./images/mask.png', uz_image.ImageType.float64)
 
     img_er_dil = img_orig.copy()
     img_dil_er = img_orig.copy()
@@ -386,19 +385,15 @@ def three_b():
 
     axs[0].imshow(original_mask, cmap='gray')
     axs[0].set(title="Original mask")
-
     axs[1].imshow(mask1, cmap='gray')
     axs[1].set(title="Using elipse")
-
     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
@@ -406,11 +401,9 @@ def three_c(image: npt.NDArray[np.float64], mask: npt.NDArray[np.uint8]):
     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)
+    img = uz_image.apply_mask_on_image(image, mask)
+
+    plt.imshow(img)
     plt.show()
 
 def three_d():
@@ -422,16 +415,13 @@ def three_d():
     would you fix that in general? (just inverting the mask if necessary doesn’t count)
     Answer:
     """
-    eagle_img_gray = uz.imread_gray('./images/eagle.jpg', uz.ImageType.uint8).astype(np.uint8)
-    eagle_img_color = uz.imread('./images/eagle.jpg', uz.ImageType.float64)
+    eagle_img_gray = uz_image.imread_gray('./images/eagle.jpg', uz_image.ImageType.uint8).astype(np.uint8)
+    eagle_img_color = uz_image.imread('./images/eagle.jpg', uz_image.ImageType.float64)
     TRESHOLD = two_e(eagle_img_gray) 
     binary_mask = eagle_img_gray.copy()
  
     binary_mask = np.where(binary_mask < TRESHOLD, 0, 1)
-    binary_mask = uz.convert_float64_array_to_uint8_array(binary_mask)
-
-    plt.imshow(binary_mask, cmap='gray')
-    plt.show()
+    binary_mask = binary_mask.astype(np.uint8)
 
     # If I would invert image here, then we would get crap 
     # So workaround:
@@ -448,7 +438,20 @@ def three_d():
     # Now invert binary mask
     binary_mask = np.where(binary_mask == 1, 0 , 1)
 
-    three_c(eagle_img_color, binary_mask)
+    removed_background_img = uz_image.apply_mask_on_image(eagle_img_color, binary_mask)
+
+    fig, axs = plt.subplots(1, 3)
+    fig.suptitle('Removing background of Eagle')
+
+    axs[0].imshow(eagle_img_color)
+    axs[0].set(title='Original image')
+    axs[1].imshow(binary_mask, cmap='gray')
+    axs[1].set(title='Computed Mask')
+    axs[2].imshow(removed_background_img, cmap='gray')
+    axs[2].set(title='Image with removed background')
+
+    plt.show() 
+
 
 def three_e():
     """
@@ -461,26 +464,24 @@ def three_e():
     from the original image (replace them with white background). Display the results
     """
     # https://stackoverflow.com/a/42812226
-    coin_img_gray = uz.imread_gray('./images/coins.jpg', uz.ImageType.uint8).astype(np.uint8)
-    coin_img_color = uz.imread('./images/coins.jpg', uz.ImageType.float64)
+    coin_img_gray = uz_image.imread_gray('./images/coins.jpg', uz_image.ImageType.uint8).astype(np.uint8)
+    coin_img_color = uz_image.imread('./images/coins.jpg', uz_image.ImageType.float64)
     TRESHOLD = two_e(coin_img_gray) 
     COIN_SIZE = 700
 
     binary_mask = coin_img_gray.copy() 
     binary_mask = np.where(binary_mask < TRESHOLD, 0, 1)
-    binary_mask = uz.convert_float64_array_to_uint8_array(binary_mask)
+    binary_mask = binary_mask.astype(np.uint8)
     
     SE_CROSS = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
     binary_mask = cv2.erode(binary_mask, SE_CROSS, iterations=3)
     binary_mask = cv2.dilate(binary_mask, SE_CROSS, iterations=2)
     SE_CROSS = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 3))
     binary_mask = cv2.dilate(binary_mask, SE_CROSS, iterations=2)
-    SE_CROSS = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 2))
-    binary_mask = cv2.dilate(binary_mask, SE_CROSS, iterations=2)
     
     # Now invert binary mask
     binary_mask = np.where(binary_mask == 1, 0 , 1)
-    binary_mask = uz.convert_float64_array_to_uint8_array(binary_mask)
+    binary_mask = binary_mask.astype(np.uint8)
 
     output = cv2.connectedComponentsWithStats(binary_mask, 4, cv2.CV_32S)
     n_blobs, im_with_separated_blobs, stats, _ = output
@@ -491,13 +492,26 @@ def three_e():
         if sizes[blob] > COIN_SIZE:
             binary_mask[im_with_separated_blobs == blob] = 0
 
-    three_c(coin_img_color, binary_mask)
+    removed_coins = uz_image.apply_mask_on_image(coin_img_color, binary_mask)
+    removed_coins[removed_coins == 0] = 1
+
+    fig, axs = plt.subplots(1, 3)
+    fig.suptitle('Removal of  >700px coins')
+
+    axs[0].imshow(coin_img_color)
+    axs[0].set(title='Original image')
+    axs[1].imshow(binary_mask, cmap='gray')
+    axs[1].set(title='Computed Mask')
+    axs[2].imshow(removed_coins)
+    axs[2].set(title='Image with coins removed')
+
+    plt.show() 
 
 
 def main() -> None:
-    #excercise_one()
+    excercise_one()
     excercise_two()
-    #excercise_three()
+    excercise_three()
 
 if __name__ == "__main__":
     main()
diff --git a/assignment1/uz_framework/image.py b/assignment1/uz_framework/image.py
index e43e583..a0db909 100644
--- a/assignment1/uz_framework/image.py
+++ b/assignment1/uz_framework/image.py
@@ -80,8 +80,6 @@ def transform_coloured_image_to_grayscale(image: npt.NDArray[np.float64]) -> npt
         for j in range(image.shape[1]):
             grayscale_image[i, j] = (image[i, j, 0] + image[i,j, 1] + image[i, j, 2]) / 3
 
-    print(grayscale_image)
-    print(grayscale_image.shape)
 
     return grayscale_image
 
@@ -213,3 +211,12 @@ def get_image_bins(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]
             counts = np.insert(counts, i - 1, 0)
 
     return counts / np.sum(counts)
+
+def apply_mask_on_image(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]], mask: npt.NDArray[np.uint8]):
+    image = image.copy()
+
+    mask = np.expand_dims(mask, axis=2) 
+    image = mask * image
+
+    return image
+