feat(utils): add functions to detect sources and generate mask

This commit adds new functions to the `rgc.utils.data` module. The `generate_mask` function is added to detect sources in an image and generate a mask. The `generate_mask_bulk` function is added to generate masks for a catalog of celestial objects. These functions are useful for image processing and analysis tasks.

The `generate_mask` function takes various parameters such as the image path, mask directory, frequency, beam size, dilation factor, and threshold values for pixel and island detection. It uses the `bdsf` library to process the image and export the generated mask.

The `generate_mask_bulk` function takes a pandas DataFrame containing the catalog of celestial objects, the image directory, mask directory, frequency, and beam size. It iterates over the catalog entries and calls the `generate_mask` function for each entry to generate masks for all the images in the catalog.

These new functions enhance the functionality of the `rgc` package and provide convenient tools for working with astronomical images.

Author: mirsazzathossain

CHANGELOG.md

@@ -14,6 +14,7 @@ classifiers=[
 dependencies = [
     "astropy>=5.2.2",
     "astroquery>=0.4.7",
+    "bdsf>=1.11.1",
     "numpy>=1.24.4",
     "pandas>=2.0.3",
     "pillow>=10.4.0",

pyproject.toml

@@ -12,6 +12,7 @@
 from pathlib import Path
 from typing import Optional, cast
 
+import bdsf
 import numpy as np
 import pandas as pd
 import torch
@@ -415,3 +416,104 @@ def remove_artifacts(folder: str, extension: list[str]) -> None:
             os.remove(os.path.join(folder, file))
 
     print(f"Artifacts removed from {folder} with extensions {', '.join(extension)}")
+
+
+def generate_mask(
+    image_path: str,
+    mask_dir: str,
+    freq: float,
+    beam: tuple[float, float, float],
+    dilation: int,
+    threshold_pixel: float = 5.0,
+    threshold_island: float = 3.0,
+) -> None:
+    """
+    Detect sources in the image and generate a mask.
+
+    :param image_path: Path to the image file
+    :type image_path: str
+
+    :param mask_dir: Path to the directory to save the mask
+    :type mask_dir: str
+
+    :param freq: Frequency of the image in MHz
+    :type freq: float
+
+    :param beam: Beam size of the image in arcsec
+    :type beam: tuple
+
+    :param dilation: Dilation factor for the mask
+    :type dilation: int
+
+    :param threshold_pixel: Threshold for island peak in number of sigma above the mean
+    :type threshold_pixel: float
+
+    :param threshold_island: Threshold for island detection in number of sigma above the mean
+    :type threshold_island: float
+    """
+    try:
+        image = bdsf.process_image(
+            image_path,
+            beam=beam,
+            thresh_isl=threshold_island,
+            thresh_pix=threshold_pixel,
+            frequency=freq,
+        )
+
+        mask_file = Path(mask_dir) / Path(image_path).name
+        Path(mask_file).parent.mkdir(parents=True, exist_ok=True)
+
+        image.export_image(
+            img_type="island_mask",
+            outfile=mask_file,
+            clobber=True,
+            mask_dilation=dilation,
+        )
+
+    except Exception:
+        print("Failed to generate mask.")
+        return None
+
+
+def generate_mask_bulk(
+    catalog: pd.DataFrame, img_dir: str, mask_dir: str, freq: float, beam: tuple[float, float, float]
+) -> None:
+    """
+    Generate masks for a catalog of celestial objects.
+
+    :param catalog: A pandas DataFrame containing the catalog of celestial objects.
+    :type catalog: pd.DataFrame
+
+    :param img_dir: The path to the directory containing the images.
+    :type img_dir: str
+
+    :param mask_dir: The path to the directory to save the masks.
+    :type mask_dir: str
+
+    :param freq: Frequency of the image in MHz
+    :type freq: float
+
+    :param beam: Beam size of the image in arcsec
+    :type beam: tuple
+    """
+    for _, entry in catalog.iterrows():
+        try:
+            filename = entry["filename"]
+            image_path = os.path.join(img_dir, f"{filename}.fits")
+            dilation = entry["dilation"]
+            threshold_pixel = entry["background sigma"]
+            threshold_island = entry["foreground sigma"]
+
+            generate_mask(
+                image_path,
+                mask_dir,
+                freq,
+                beam,
+                dilation,
+                threshold_pixel,
+                threshold_island,
+            )
+
+        except Exception as err:
+            print(f"Failed to generate mask. {err}")
+            return None

rgc/utils/data.py

@@ -0,0 +1,126 @@
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from rgc.utils.data import generate_mask
+
+
+class TestGenerateMask(unittest.TestCase):
+    @patch("rgc.utils.data.bdsf.process_image", return_value=MagicMock())
+    @patch("rgc.utils.data.Path.mkdir")
+    @patch("rgc.utils.data.Path.name", new_callable=MagicMock(return_value="image.fits"))
+    @patch("rgc.utils.data.Path", autospec=True)
+    @patch("rgc.utils.data.print")
+    def test_generate_mask_success(self, mock_print, mock_path, mock_name, mock_mkdir, mock_process_image):
+        # Setup mocks
+        mock_path_instance = mock_path.return_value
+        # Ensure that the truediv operator returns the correct Path
+        mock_path_instance.__truediv__.return_value = Path("path/to/mask_dir/image.fits")
+        mock_path_instance.parent.mkdir.return_value = None
+        mock_path_instance.name = "image.fits"  # Correct the name attribute
+
+        # Setup test parameters
+        image_path = "path/to/image.fits"
+        mask_dir = "path/to/mask_dir"
+        freq = 1400.0
+        beam = (5.0, 5.0, 5.0)
+        dilation = 2
+        threshold_pixel = 5.0
+        threshold_island = 3.0
+
+        # Run the function
+        generate_mask(
+            image_path=image_path,
+            mask_dir=mask_dir,
+            freq=freq,
+            beam=beam,
+            dilation=dilation,
+            threshold_pixel=threshold_pixel,
+            threshold_island=threshold_island,
+        )
+
+        # Verify that print was not called
+        mock_print.assert_not_called()
+
+        # Verify that the mask file path was generated correctly
+        expected_mask_file = Path(mask_dir) / "image.fits"
+        mock_path_instance.__truediv__.assert_called_once_with("image.fits")
+
+        # Verify that the parent directory was created
+        mock_path_instance.parent.mkdir.assert_called_once_with(parents=True, exist_ok=True)
+
+        # Verify that process_image was called correctly
+        mock_process_image.assert_called_once_with(
+            image_path, beam=beam, thresh_isl=threshold_island, thresh_pix=threshold_pixel, frequency=freq
+        )
+
+        # Verify that export_image is called on the mock returned by process_image
+        mock_process_image.return_value.export_image.assert_called_once_with(
+            img_type="island_mask",
+            outfile=expected_mask_file,
+            clobber=True,
+            mask_dilation=dilation,
+        )
+
+    @patch("rgc.utils.data.bdsf.process_image", return_value=MagicMock())
+    @patch("rgc.utils.data.Path.mkdir", side_effect=PermissionError("Permission denied"))
+    @patch("rgc.utils.data.Path.parent", new_callable=MagicMock)
+    @patch("rgc.utils.data.Path.name", new_callable=MagicMock(return_value="image.fits"))
+    @patch("rgc.utils.data.Path", autospec=True)
+    @patch("rgc.utils.data.print")
+    def test_generate_mask_permission_error(
+        self, mock_print, mock_path, mock_name, mock_parent, mock_mkdir, mock_process_image
+    ):
+        # Setup mocks
+        mock_path.return_value = Path("/mock/path")
+        mock_parent.return_value = mock_path
+        mock_process_image.return_value = MagicMock()
+
+        # Setup test parameters
+        image_path = "path/to/image.fits"
+        mask_dir = "path/to/mask_dir"
+        freq = 1400.0
+        beam = (5.0, 5.0, 5.0)
+        dilation = 2
+        threshold_pixel = 5.0
+        threshold_island = 3.0
+
+        # Run the function
+        generate_mask(
+            image_path=image_path,
+            mask_dir=mask_dir,
+            freq=freq,
+            beam=beam,
+            dilation=dilation,
+            threshold_pixel=threshold_pixel,
+            threshold_island=threshold_island,
+        )
+
+        # Verify that print was called with the correct error message
+        mock_print.assert_called_once_with("Failed to generate mask.")
+
+    @patch("rgc.utils.data.bdsf.process_image", side_effect=Exception("Process failed"))
+    @patch("rgc.utils.data.print")
+    def test_generate_mask_failure(self, mock_print, mock_process_image):
+        # Setup test parameters
+        image_path = "path/to/image.fits"
+        mask_dir = "path/to/mask_dir"
+        freq = 1400.0
+        beam = (5.0, 5.0, 5.0)
+        dilation = 2
+        threshold_pixel = 5.0
+        threshold_island = 3.0
+
+        # Run the function
+        generate_mask(
+            image_path=image_path,
+            mask_dir=mask_dir,
+            freq=freq,
+            beam=beam,
+            dilation=dilation,
+            threshold_pixel=threshold_pixel,
+            threshold_island=threshold_island,
+        )
+
+        # Verify that print was called with the correct error message
+        mock_print.assert_called_once_with("Failed to generate mask.")

tests/test_generate_mask.py

@@ -0,0 +1,104 @@
+import unittest
+from unittest.mock import call, patch
+
+import pandas as pd
+
+from rgc.utils.data import generate_mask_bulk
+
+
+class TestGenerateMaskBulk(unittest.TestCase):
+    @patch("rgc.utils.data.generate_mask")
+    @patch("rgc.utils.data.os.path.join", return_value="mock_image_path.fits")
+    @patch("rgc.utils.data.print")  # Mock print to suppress output during the test
+    def test_generate_mask_bulk_success(self, mock_print, mock_join, mock_generate_mask):
+        # Create a mock catalog DataFrame
+        data = {
+            "filename": ["image1", "image2"],
+            "dilation": [2, 3],
+            "background sigma": [5.0, 6.0],
+            "foreground sigma": [3.0, 4.0],
+        }
+        catalog = pd.DataFrame(data)
+
+        # Test parameters
+        img_dir = "mock/img/dir"
+        mask_dir = "mock/mask/dir"
+        freq = 1400.0
+        beam = (5.0, 5.0, 5.0)
+
+        # Run the function
+        generate_mask_bulk(
+            catalog=catalog,
+            img_dir=img_dir,
+            mask_dir=mask_dir,
+            freq=freq,
+            beam=beam,
+        )
+
+        # Check that os.path.join was called with the correct parameters
+        expected_calls = [call(img_dir, "image1.fits"), call(img_dir, "image2.fits")]
+        mock_join.assert_has_calls(expected_calls, any_order=False)
+
+        # Check that generate_mask was called twice with the correct parameters
+        expected_mask_calls = [
+            call(
+                "mock_image_path.fits",
+                mask_dir,
+                freq,
+                beam,
+                2,  # dilation for image1
+                5.0,  # background sigma for image1
+                3.0,  # foreground sigma for image1
+            ),
+            call(
+                "mock_image_path.fits",
+                mask_dir,
+                freq,
+                beam,
+                3,  # dilation for image2
+                6.0,  # background sigma for image2
+                4.0,  # foreground sigma for image2
+            ),
+        ]
+        mock_generate_mask.assert_has_calls(expected_mask_calls, any_order=False)
+
+        # Ensure no errors were printed
+        mock_print.assert_not_called()
+
+    @patch("rgc.utils.data.generate_mask", side_effect=Exception("Mocked error"))
+    @patch("rgc.utils.data.os.path.join", return_value="mock_image_path.fits")
+    @patch("rgc.utils.data.print")
+    def test_generate_mask_bulk_failure(self, mock_print, mock_join, mock_generate_mask):
+        # Create a mock catalog DataFrame
+        data = {
+            "filename": ["image1"],
+            "dilation": [2],
+            "background sigma": [5.0],
+            "foreground sigma": [3.0],
+        }
+        catalog = pd.DataFrame(data)
+
+        # Test parameters
+        img_dir = "mock/img/dir"
+        mask_dir = "mock/mask/dir"
+        freq = 1400.0
+        beam = (5.0, 5.0, 5.0)
+
+        # Run the function
+        generate_mask_bulk(
+            catalog=catalog,
+            img_dir=img_dir,
+            mask_dir=mask_dir,
+            freq=freq,
+            beam=beam,
+        )
+
+        # Verify that the error message was printed
+        mock_print.assert_called_once_with("Failed to generate mask. Mocked error")
+
+        # Check that generate_mask was called once before the exception
+        mock_generate_mask.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()