Merge pull request #15 from mirsazzathossain/dev

feat: Add Masking Functionality for Single and Batch Image Processing

Author: mirsazzathossain

.github/ISSUE_TEMPLATE/bug_report.md

@@ -1,17 +1,17 @@
 ---
 name: Bug report
 about: Create a report to help us improve
-title: ''
-labels: ''
-assignees: ''
-
+title: ""
+labels: ""
+assignees: ""
 ---
 
 **Describe the bug**
 A clear and concise description of what the bug is.
 
 **To Reproduce**
 Steps to reproduce the behavior:
+
 1. Go to '...'
 2. Click on '....'
 3. Scroll down to '....'
@@ -24,15 +24,17 @@ A clear and concise description of what you expected to happen.
 If applicable, add screenshots to help explain your problem.
 
 **Desktop (please complete the following information):**
- - OS: [e.g. iOS]
- - Browser [e.g. chrome, safari]
- - Version [e.g. 22]
+
+- OS: [e.g. iOS]
+- Browser [e.g. chrome, safari]
+- Version [e.g. 22]
 
 **Smartphone (please complete the following information):**
- - Device: [e.g. iPhone6]
- - OS: [e.g. iOS8.1]
- - Browser [e.g. stock browser, safari]
- - Version [e.g. 22]
+
+- Device: [e.g. iPhone6]
+- OS: [e.g. iOS8.1]
+- Browser [e.g. stock browser, safari]
+- Version [e.g. 22]
 
 **Additional context**
 Add any other context about the problem here.

.github/ISSUE_TEMPLATE/feature_request.md

@@ -1,10 +1,9 @@
 ---
 name: Feature request
 about: Suggest an idea for this project
-title: ''
-labels: ''
-assignees: ''
-
+title: ""
+labels: ""
+assignees: ""
 ---
 
 **Is your feature request related to a problem? Please describe.**

CHANGELOG.md

@@ -124,11 +124,11 @@ def __init__(self) -> None:
 
 class _FileNotFoundError(Exception):
     """
-    An exception to be raised when the FITS file is not found.
+    An exception to be raised when a file is not found.
     """
 
-    def __init__(self, fits_file: str) -> None:
-        super().__init__(f"File {fits_file} not found.")
+    def __init__(self, message: str = "File not found.") -> None:
+        super().__init__(message)
 
 
 def fits_to_png(fits_file: str, img_size: Optional[tuple[int, int]] = None) -> Image.Image:
@@ -189,3 +189,77 @@ def fits_to_png_bulk(fits_dir: str, png_dir: str, img_size: Optional[tuple[int,
 
         if image is not None:
             image.save(png_file)
+
+
+def mask_image(image: Image.Image, mask: Image.Image) -> Image.Image:
+    """
+    Mask an image with a given mask image.
+
+    :param image: The image to be masked.
+    :type image: Image.Image
+
+    :param mask: The mask image.
+    :type mask: Image.Image
+
+    :return: A PIL Image object containing the masked image.
+    :rtype: Image.Image
+    """
+    image_array = np.array(image)
+    mask_array = np.array(mask)
+
+    if image_array.shape != mask_array.shape:
+        raise _ImageMaskDimensionError()
+
+    masked_array = np.where(mask_array == 0, 0, image_array)
+    masked_image = Image.fromarray(masked_array, mode="L")
+
+    return cast(Image.Image, masked_image)
+
+
+class _ImageMaskDimensionError(Exception):
+    """
+    An exception to be raised when the dimensions of the image and mask do not match.
+    """
+
+    def __init__(self) -> None:
+        super().__init__("Image and mask must have the same dimensions.")
+
+
+class _ImageMaskCountMismatchError(Exception):
+    """
+    An exception to be raised when the number of images and masks do not match.
+    """
+
+    def __init__(self, message: str = "Number of images and masks must match and be non-zero.") -> None:
+        super().__init__(message)
+
+
+def mask_image_bulk(image_dir: str, mask_dir: str, masked_dir: str) -> None:
+    image_paths = sorted(Path(image_dir).glob("*.png"))
+    mask_paths = sorted(Path(mask_dir).glob("*.png"))
+
+    if len(image_paths) == 0 or len(mask_paths) == 0:
+        raise _FileNotFoundError()
+
+    if len(image_paths) != len(mask_paths):
+        raise _ImageMaskCountMismatchError() from None
+
+    os.makedirs(masked_dir, exist_ok=True)
+
+    for image_path in image_paths:
+        mask_path = Path(mask_dir) / image_path.name
+
+        if not mask_path.exists():
+            print(f"Skipping {image_path.name} due to missing mask.")
+            continue
+
+        image = Image.open(image_path)
+        mask = Image.open(mask_path)
+
+        if image.size != mask.size:
+            print(f"Skipping {image_path.name} due to mismatched dimensions.")
+            continue
+        else:
+            masked_image = mask_image(image, mask)
+
+        masked_image.save(Path(masked_dir) / image_path.name)

rgc/utils/data.py

@@ -0,0 +1,64 @@
+import unittest
+
+import numpy as np
+from PIL import Image
+
+from rgc.utils.data import _ImageMaskDimensionError, mask_image
+
+
+class TestMaskImage(unittest.TestCase):
+    def setUp(self):
+        # Create sample images and masks for testing
+        self.image_array = np.array([[100, 150, 200], [50, 75, 100], [0, 25, 50]], dtype=np.uint8)
+        self.mask_array = np.array([[1, 0, 1], [0, 1, 0], [1, 1, 0]], dtype=np.uint8)
+        self.image = Image.fromarray(self.image_array, mode="L")
+        self.mask = Image.fromarray(self.mask_array, mode="L")
+
+    def test_mask_all_zeros(self):
+        zero_mask_array = np.zeros_like(self.mask_array)
+        zero_mask = Image.fromarray(zero_mask_array, mode="L")
+
+        expected_array = np.zeros_like(self.image_array)
+        _ = Image.fromarray(expected_array, mode="L")
+
+        result_image = mask_image(self.image, zero_mask)
+        result_array = np.array(result_image)
+
+        np.testing.assert_array_equal(result_array, expected_array)
+
+    def test_mask_all_ones(self):
+        ones_mask_array = np.ones_like(self.mask_array)
+        ones_mask = Image.fromarray(ones_mask_array, mode="L")
+
+        expected_array = self.image_array.copy()
+        _ = Image.fromarray(expected_array, mode="L")
+
+        result_image = mask_image(self.image, ones_mask)
+        result_array = np.array(result_image)
+
+        np.testing.assert_array_equal(result_array, expected_array)
+
+    def test_non_matching_dimension(self):
+        small_mask_array = np.array([[1, 0]], dtype=np.uint8)
+        small_mask = Image.fromarray(small_mask_array, mode="L")
+
+        with self.assertRaises(_ImageMaskDimensionError):
+            mask_image(self.image, small_mask)
+
+    def test_empty_image(self):
+        empty_image_array = np.array([[]], dtype=np.uint8)
+        empty_image = Image.fromarray(empty_image_array, mode="L")
+
+        with self.assertRaises(_ImageMaskDimensionError):
+            mask_image(empty_image, self.mask)
+
+    def test_empty_mask(self):
+        empty_mask_array = np.array([[]], dtype=np.uint8)
+        empty_mask = Image.fromarray(empty_mask_array, mode="L")
+
+        with self.assertRaises(_ImageMaskDimensionError):
+            mask_image(self.image, empty_mask)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_mask_image.py

@@ -0,0 +1,120 @@
+import os
+import shutil
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+from PIL import Image
+
+from rgc.utils.data import _FileNotFoundError, _ImageMaskCountMismatchError, mask_image_bulk
+
+
+class TestMaskImageBulk(unittest.TestCase):
+    def setUp(self):
+        self.image_dir = tempfile.mkdtemp()
+        self.mask_dir = tempfile.mkdtemp()
+        self.masked_dir = tempfile.mkdtemp()
+
+        self.image_array = np.array([[100, 150, 200], [50, 75, 100], [0, 25, 50]], dtype=np.uint8)
+        self.mask_array = np.array([[1, 0, 1], [0, 1, 0], [1, 1, 0]], dtype=np.uint8)
+
+        self.image_path = Path(self.image_dir) / "test_image.png"
+        self.mask_path = Path(self.mask_dir) / "test_image.png"
+
+        Image.fromarray(self.image_array, mode="L").save(self.image_path)
+        Image.fromarray(self.mask_array, mode="L").save(self.mask_path)
+
+    def tearDown(self):
+        shutil.rmtree(self.image_dir)
+        shutil.rmtree(self.mask_dir)
+        shutil.rmtree(self.masked_dir)
+
+    def test_mask_image_bulk(self):
+        mask_image_bulk(self.image_dir, self.mask_dir, self.masked_dir)
+        masked_file_path = Path(self.masked_dir) / "test_image.png"
+        self.assertTrue(masked_file_path.exists())
+        masked_image = Image.open(masked_file_path)
+        masked_array = np.array(masked_image)
+        expected_array = np.array([[100, 0, 200], [0, 75, 0], [0, 25, 0]], dtype=np.uint8)
+        np.testing.assert_array_equal(masked_array, expected_array)
+
+    @patch("builtins.print")
+    def test_dimension_mismatch(self, mock_print):
+        # Ensure mask_dir is empty
+        for mask_file in Path(self.mask_dir).glob("*.png"):
+            os.remove(mask_file)
+
+        # Create a mask with a different dimension
+        mismatch_mask_array = np.array([[1, 0]], dtype=np.uint8)
+        mismatch_mask_path = Path(self.mask_dir) / "test_image.png"
+        Image.fromarray(mismatch_mask_array, mode="L").save(mismatch_mask_path)
+
+        # Ensure image_dir contains only the test image
+        for image_file in Path(self.image_dir).glob("*.png"):
+            os.remove(image_file)
+
+        Image.fromarray(self.image_array, mode="L").save(Path(self.image_dir) / "test_image.png")
+
+        # Run the function and check if the dimension mismatch is handled
+        mask_image_bulk(self.image_dir, self.mask_dir, self.masked_dir)
+
+        # Check if masked directory is still empty
+        self.assertFalse(
+            list(Path(self.masked_dir).glob("*.png")),
+            "Masked directory should be empty if there is a dimension mismatch",
+        )
+
+        # Verify that the print statement was made
+        # Ensure to check the exact message your code prints
+        mock_print.assert_called_with("Skipping test_image.png due to mismatched dimensions.")
+
+    def test_missing_mask_file(self):
+        # Create a directory with an image but without a corresponding mask
+        missing_mask_dir = tempfile.mkdtemp()
+        Image.fromarray(self.image_array, mode="L").save(Path(missing_mask_dir) / "fake_image.png")
+
+        mask_image_bulk(self.image_dir, missing_mask_dir, self.masked_dir)
+
+        # Check that masked directory is still empty
+        self.assertFalse(os.listdir(self.masked_dir), "Masked directory should be empty if mask file is missing")
+
+        shutil.rmtree(missing_mask_dir)
+
+    def test_empty_image_dir(self):
+        empty_image_dir = tempfile.mkdtemp()
+        with self.assertRaises(_FileNotFoundError):
+            mask_image_bulk(empty_image_dir, self.mask_dir, self.masked_dir)
+        shutil.rmtree(empty_image_dir)
+
+    def test_empty_mask_dir(self):
+        empty_mask_dir = tempfile.mkdtemp()
+        with self.assertRaises(_FileNotFoundError):
+            mask_image_bulk(self.image_dir, empty_mask_dir, self.masked_dir)
+        shutil.rmtree(empty_mask_dir)
+
+    def test_non_matching_images_and_masks(self):
+        extra_image_dir = tempfile.mkdtemp()
+        extra_mask_dir = tempfile.mkdtemp()
+
+        extra_image_path = Path(extra_image_dir) / "extra_image.png"
+        Image.fromarray(self.image_array, mode="L").save(extra_image_path)
+
+        extra_image_path = Path(extra_image_dir) / "extra_image_2.png"
+        Image.fromarray(self.image_array, mode="L").save(extra_image_path)
+
+        with self.assertRaises(_ImageMaskCountMismatchError):
+            mask_image_bulk(extra_image_dir, self.mask_dir, self.masked_dir)
+
+        extra_mask_path = Path(extra_mask_dir) / "extra_mask.png"
+        Image.fromarray(self.mask_array, mode="L").save(extra_mask_path)
+
+        self.assertFalse(os.listdir(self.masked_dir))
+
+        shutil.rmtree(extra_image_dir)
+        shutil.rmtree(extra_mask_dir)
+
+
+if __name__ == "__main__":
+    unittest.main()