Extracted background removal from daophot and removebackground into i…

…ts own class

pyobs/images/processors/_daobackgroundremover.py

@@ -0,0 +1,37 @@
+from typing import Tuple
+
+import numpy as np
+from astropy.stats import SigmaClip
+
+from pyobs.images import Image
+
+
+class _DaoBackgroundRemover:
+    def __init__(self, sigma: float, box_size: Tuple[int, int], filter_size: Tuple[int, int]):
+        from photutils.background import MedianBackground
+
+        self._sigma_clip = SigmaClip(sigma=sigma)
+        self._box_size = box_size
+        self._filter_size = filter_size
+
+        self._bkg_estimator = MedianBackground()
+
+    def __call__(self, image: Image) -> Image:
+        background = self._estimate_background(image)
+        return self._remove_background(image, background)
+
+    def _estimate_background(self, image: Image):
+        from photutils.background import Background2D
+
+        bkg = Background2D(
+            image.data, box_size=self._box_size, filter_size=self._filter_size, sigma_clip=self._sigma_clip,
+            bkg_estimator=self._bkg_estimator, mask=image.mask
+        )
+
+        return bkg.background
+
+    @staticmethod
+    def _remove_background(image: Image, background: np.ndarray) -> Image:
+        output_image = image.copy()
+        output_image.data = output_image.data - background
+        return output_image

pyobs/images/processors/detection/daophot.py

@@ -3,12 +3,13 @@
 from typing import Tuple, Any
 
 import numpy as np
+from astropy.stats import sigma_clipped_stats
 from astropy.table import Table
-from astropy.stats import SigmaClip, sigma_clipped_stats
 
 from pyobs.images import Image
 from ._source_catalog import _SourceCatalog
 from .sourcedetection import SourceDetection
+from .._daobackgroundremover import _DaoBackgroundRemover
 
 log = logging.getLogger(__name__)
 
@@ -43,29 +44,8 @@ def __init__(
         # store
         self.fwhm = fwhm
         self.threshold = threshold
-        self.bkg_sigma = bkg_sigma
-        self.bkg_box_size = bkg_box_size
-        self.bkg_filter_size = bkg_filter_size
-
-    def _estimate_background(self, data: np.ndarray, mask: np.ndarray) -> np.ndarray:
-        from photutils import Background2D, MedianBackground
-
-        sigma_clip = SigmaClip(sigma=self.bkg_sigma)
-        bkg_estimator = MedianBackground()
-        bkg = Background2D(
-            data,
-            self.bkg_box_size,
-            filter_size=self.bkg_filter_size,
-            sigma_clip=sigma_clip,
-            bkg_estimator=bkg_estimator,
-            mask=mask,
-        )
-
-        return bkg.background
-
-    def _remove_background_from_data(self, data, mask) -> np.ndarray:
-        background = self._estimate_background(data, mask)
-        return data - background
+
+        self._background_remover = _DaoBackgroundRemover(bkg_sigma, bkg_box_size, bkg_filter_size)
 
     async def _find_stars(self, data: np.ndarray, std: int) -> Table:
         from photutils import DAOStarFinder
@@ -88,9 +68,9 @@ async def __call__(self, image: Image) -> Image:
         if image.data is None:
             log.warning("No data found in image.")
             return image
-        image_data = image.data.astype(float)
 
-        background_corrected_data = self._remove_background_from_data(image_data, image.mask)
+        background_corrected_image = self._background_remover(image)
+        background_corrected_data = background_corrected_image.data.astype(float)
 
         _, median, std = sigma_clipped_stats(background_corrected_data, sigma=3.0)
 

pyobs/images/processors/misc/removebackground.py

@@ -4,6 +4,7 @@
 
 from pyobs.images.processor import ImageProcessor
 from pyobs.images import Image
+from pyobs.images.processors._daobackgroundremover import _DaoBackgroundRemover
 
 log = logging.getLogger(__name__)
 
@@ -29,10 +30,7 @@ def __init__(
         """
         ImageProcessor.__init__(self, **kwargs)
 
-        # store
-        self.sigma = sigma
-        self.box_size = box_size
-        self.filter_size = filter_size
+        self._background_remover = _DaoBackgroundRemover(sigma, box_size, filter_size)
 
     async def __call__(self, image: Image) -> Image:
         """Remove background from image.
@@ -43,21 +41,9 @@ async def __call__(self, image: Image) -> Image:
         Returns:
             Image without background.
         """
-        from photutils.background import Background2D, MedianBackground
 
-        # init objects
-        sigma_clip = SigmaClip(sigma=self.sigma)
-        bkg_estimator = MedianBackground()
+        return self._background_remover(image)
 
-        # calculate background
-        bkg = Background2D(
-            image.data, self.box_size, filter_size=self.filter_size, sigma_clip=sigma_clip, bkg_estimator=bkg_estimator
-        )
-
-        # copy image and remove background
-        img = image.copy()
-        img.data = img.data - bkg.background
-        return img
 
 
 __all__ = ["RemoveBackground"]

tests/images/processors/detection/test_daophot.py

@@ -10,9 +10,9 @@ def test_init_default():
 
     assert detector.fwhm == 3.0
     assert detector.threshold == 4.0
-    assert detector.bkg_sigma == 3.0
-    assert detector.bkg_box_size == (50, 50)
-    assert detector.bkg_filter_size == (3, 3)
+    assert detector._background_remover._sigma_clip.sigma == 3.0
+    assert detector._background_remover._box_size == (50, 50)
+    assert detector._background_remover._filter_size == (3, 3)
 
 
 def test_init():
@@ -26,9 +26,9 @@ def test_init():
 
     assert detector.fwhm == fwhm
     assert detector.threshold == threshold
-    assert detector.bkg_sigma == bkg_sigma
-    assert detector.bkg_box_size == bkg_box_size
-    assert detector.bkg_filter_size == bkg_filter_size
+    assert detector._background_remover._sigma_clip.sigma == bkg_sigma
+    assert detector._background_remover._box_size == bkg_box_size
+    assert detector._background_remover._filter_size == bkg_filter_size
 
 
 @pytest.mark.asyncio

tests/images/processors/misc/test_removebackground.py

@@ -1,3 +1,5 @@
+from unittest.mock import Mock
+
 import numpy as np
 import photutils.background
 import pytest
@@ -12,28 +14,30 @@ def test_init():
     filter_size = (3, 3)
     remover = RemoveBackground(sigma, box_size, filter_size)
 
-    assert remover.sigma == sigma
-    assert remover.box_size == box_size
-    assert remover.filter_size == filter_size
+    assert remover._background_remover._sigma_clip.sigma == sigma
+    assert remover._background_remover._box_size == box_size
+    assert remover._background_remover._filter_size == filter_size
 
 
 def test_init_default():
     remover = RemoveBackground()
-    assert remover.sigma == 3.0
-    assert remover.box_size == (50, 50)
-    assert remover.filter_size == (3, 3)
+    assert remover._background_remover._sigma_clip.sigma == 3.0
+    assert remover._background_remover._box_size == (50, 50)
+    assert remover._background_remover._filter_size == (3, 3)
 
 
 @pytest.mark.asyncio
-async def test_call_const_background(mocker):
+async def test_call_const_background():
     sigma = 3.0
     box_size = (1, 1)
     filter_size = (3, 3)
-    spy = mocker.spy(photutils.background.Background2D, "__init__")
     remover = RemoveBackground(sigma, box_size, filter_size)
 
     image = Image(data=np.ones((20, 20)))
-    output_image = await remover(image)
+    output_image = Image(data=np.ones((20, 20) * 2))
+    remover._background_remover = Mock(return_value=output_image)
+
+    result = await remover(image)
 
-    spy.assert_called_once()
-    np.testing.assert_array_equal(output_image.data, np.zeros((20, 20)))
+    np.testing.assert_array_equal(result.data, output_image.data)
+    remover._background_remover.assert_called_once()

tests/images/processors/test_removebackground.py

@@ -0,0 +1,52 @@
+import numpy as np
+import photutils.background
+
+from pyobs.images import Image
+from pyobs.images.processors._daobackgroundremover import _DaoBackgroundRemover
+
+
+def test_init():
+    sigma = 1.0
+    box_size = (10, 10)
+    filter_size = (3, 3)
+    remover = _DaoBackgroundRemover(sigma, box_size, filter_size)
+
+    assert remover._sigma_clip.sigma == sigma
+    assert remover._box_size == box_size
+    assert remover._filter_size == filter_size
+
+
+def test_estimate_background_background2d_call(mocker):
+    sigma = 3.0
+    box_size = (1, 1)
+    filter_size = (3, 3)
+    spy = mocker.spy(photutils.background.Background2D, "__init__")
+    remover = _DaoBackgroundRemover(sigma, box_size, filter_size)
+
+    data = np.ones((20, 20))
+    mask = np.zeros((20, 20)).astype(bool)
+    image = Image(data=data, mask=mask)
+    remover._estimate_background(image)
+    spy.assert_called_once()
+
+    args = spy.call_args[0]
+    np.testing.assert_array_equal(args[1], data)
+    kwargs = spy.call_args.kwargs
+    assert kwargs["box_size"] == box_size
+    assert kwargs["filter_size"] == filter_size
+    assert kwargs["sigma_clip"].sigma == sigma
+    assert kwargs["bkg_estimator"] == remover._bkg_estimator
+    np.testing.assert_array_equal(kwargs["mask"], mask)
+
+
+def test_call_const_background(mocker):
+    sigma = 3.0
+    box_size = (1, 1)
+    filter_size = (3, 3)
+
+    remover = _DaoBackgroundRemover(sigma, box_size, filter_size)
+
+    image = Image(data=np.ones((20, 20)))
+    output_image = remover(image)
+
+    np.testing.assert_array_equal(output_image.data, np.zeros((20, 20)))