Add double alpha channel support and improve metadata behaviours for BackgroundCompositor

satpy/composites/__init__.py

@@ -1,4 +1,4 @@
 # Copyright (c) 2015-2023 Satpy developers
 #
 # This file is part of satpy.
 #
@@ -1670,25 +1670,56 @@
 
 
 class BackgroundCompositor(GenericCompositor):
-    """A compositor that overlays one composite on top of another."""
+    """A compositor that overlays one composite on top of another.
+
+    The output image mode will be determined by both foreground and background. Generally, when the background has
+    an alpha band, the output image will also have one.
+        # L/L -> L
+        # L/LA -> LA
+        # L/RGB -> RGB
+        # L/RGBA -> RGBA
+        # LA/L -> L
+        # LA/LA -> LA
+        # LA/RGB -> RGB
+        # LA/RGBA -> RGBA
+        # RGB/L -> RGB
+        # RGB/LA -> RGBA
+        # RGB/RGB -> RGB
+        # RGB/RGBA -> RGBA
+        # RGBA/L -> RGB
+        # RGBA/LA -> RGBA
+        # RGBA/RGB -> RGB
+        # RGBA/RGBA -> RGBA
+
+    """
 
     def __call__(self, projectables, *args, **kwargs):
         """Call the compositor."""
         projectables = self.match_data_arrays(projectables)
+
         # Get enhanced datasets
         foreground = enhance2dataset(projectables[0], convert_p=True)
         background = enhance2dataset(projectables[1], convert_p=True)
-        # Adjust bands so that they match
-        # L/RGB -> RGB/RGB
-        # LA/RGB -> RGBA/RGBA
-        # RGB/RGBA -> RGBA/RGBA
+        before_bg_mode = background.attrs["mode"]
+
+        # Adjust bands so that they have the same mode
         foreground = add_bands(foreground, background["bands"])
         background = add_bands(background, foreground["bands"])
 
+        # It's important whether the alpha band of background is initially generated, e.g. by CloudCompositor
+        # The result will be used to determine the output image mode
+        initial_bg_alpha = "A" in before_bg_mode
+
         attrs = self._combine_metadata_with_mode_and_sensor(foreground, background)
-        data = self._get_merged_image_data(foreground, background)
+        if "A" not in foreground.attrs["mode"] and "A" not in background.attrs["mode"]:
+            data = self._simple_overlay(foreground, background)
+        else:
+            data = self._get_merged_image_data(foreground, background, initial_bg_alpha=initial_bg_alpha)
+        for data_arr in data:
+            data_arr.attrs = attrs
         res = super(BackgroundCompositor, self).__call__(data, **kwargs)
-        res.attrs.update(attrs)
+        attrs.update(res.attrs)
+        res.attrs = attrs
         return res
 
     def _combine_metadata_with_mode_and_sensor(self,
@@ -1707,27 +1738,61 @@
 
     @staticmethod
     def _get_merged_image_data(foreground: xr.DataArray,
-                               background: xr.DataArray
+                               background: xr.DataArray,
+                               initial_bg_alpha: bool,
                                ) -> list[xr.DataArray]:
-        if "A" in foreground.attrs["mode"]:
-            # Use alpha channel as weight and blend the two composites
-            alpha = foreground.sel(bands="A")
-            data = []
-            # NOTE: there's no alpha band in the output image, it will
-            # be added by the data writer
-            for band in foreground.mode[:-1]:
-                fg_band = foreground.sel(bands=band)
-                bg_band = background.sel(bands=band)
-                chan = (fg_band * alpha + bg_band * (1 - alpha))
-                chan = xr.where(chan.isnull(), bg_band, chan)
-                data.append(chan)
-        else:
-            data_arr = xr.where(foreground.isnull(), background, foreground)
-            # Split to separate bands so the mode is correct
-            data = [data_arr.sel(bands=b) for b in data_arr["bands"]]
+        # For more info about alpha compositing please review https://en.wikipedia.org/wiki/Alpha_compositing
+        alpha_fore = _get_alpha(foreground)
+        alpha_back = _get_alpha(background)
+        new_alpha = alpha_fore + alpha_back * (1 - alpha_fore)
+
+        data = []
+
+        # Pass the image data (alpha band will be dropped temporally) to the writer
+        output_mode = background.attrs["mode"].replace("A", "")
+
+        for band in output_mode:
+            fg_band = foreground.sel(bands=band)
+            bg_band = background.sel(bands=band)
+            # Do the alpha compositing
+            chan = (fg_band * alpha_fore + bg_band * alpha_back * (1 - alpha_fore)) / new_alpha
+            # Fill the NaN area with background
+            chan = xr.where(chan.isnull(), bg_band * alpha_back, chan)
+            chan["bands"] = band
+            data.append(chan)
+
+        # If background has an initial alpha band, it will also be passed to the writer
+        if initial_bg_alpha:
+            new_alpha["bands"] = "A"
+            data.append(new_alpha)
 
         return data
 
+    @staticmethod
+    def _simple_overlay(foreground: xr.DataArray,
+                        background: xr.DataArray,) -> list[xr.DataArray]:
+        # This is for the case when no alpha bands are involved
+        # Just simply lay the foreground upon background
+        data_arr = xr.where(foreground.isnull(), background, foreground)
+        # Split to separate bands so the mode is correct
+        data = [data_arr.sel(bands=b) for b in data_arr["bands"]]
+
+        return data
+
+
+def _get_alpha(dataset: xr.DataArray):
+    # 1. This function is only used by _get_merged_image_data
+    # 2. Both foreground and background have been through add_bands, so they have the same mode
+    # 3. If none of them has alpha band, they will be passed to _simple_overlay not _get_merged_image_data
+    # So any dataset(whether foreground or background) passed to this function has an alpha band for certain
+    # We will use it directly
+    alpha = dataset.sel(bands="A")
+    # There could be NaNs in the alpha
+    # Replace them with 0 to prevent cases like 1 + nan = nan, so they won't affect new_alpha
+    alpha = xr.where(alpha.isnull(), 0, alpha)
+
+    return alpha
+
 
 class MaskingCompositor(GenericCompositor):
     """A compositor that masks e.g. IR 10.8 channel data using cloud products from NWC SAF."""

satpy/tests/reader_tests/test_olci_nc.py

@@ -335,10 +335,10 @@ def test_bitflags_with_dataarray_without_flags(self):
                  "CLOUD_MARGIN", "CLOUD_AMBIGUOUS", "LOWRW", "LAND"]
 
         mask = reduce(np.logical_or, [bflags[item] for item in items])
-        expected = np.array([True, False,  True,  True,  True,  True, False,
-                             False,  True, True, False, False, False, False,
-                             False, False, False,  True, False,  True, False,
-                             False, False,  True,  True, False, False, True,
+        expected = np.array([True, False, True, True, True, True, False,
+                             False, True, True, False, False, False, False,
+                             False, False, False, True, False, True, False,
+                             False, False, True, True, False, False, True,
                              False])
         assert all(mask == expected)
 
@@ -367,9 +367,9 @@ def test_bitflags_with_custom_flag_list(self):
                  "CLOUD_MARGIN", "CLOUD_AMBIGUOUS", "LOWRW", "LAND"]
 
         mask = reduce(np.logical_or, [bflags[item] for item in items])
-        expected = np.array([True, False,  True,  True,  True,  True, False,
-                             False,  True, True, False, False, False, False,
-                             False, False, False,  True, False,  True, False,
-                             False, False,  True,  True, False, False, True,
+        expected = np.array([True, False, True, True, True, True, False,
+                             False, True, True, False, False, False, False,
+                             False, False, False, True, False, True, False,
+                             False, False, True, True, False, False, True,
                              False])
         assert all(mask == expected)

satpy/tests/test_composites.py

@@ -1484,31 +1484,52 @@ def setup_class(cls):
                 [[1., 0.5], [0., np.nan]],
                 [[1., 0.5], [0., np.nan]]]),
             "RGBA": np.array([
-                [[1.0, 0.5], [0.0, np.nan]],
-                [[1.0, 0.5], [0.0, np.nan]],
-                [[1.0, 0.5], [0.0, np.nan]],
-                [[0.5, 0.5], [0.5, 0.5]]]),
+                [[1., 0.5], [0., np.nan]],
+                [[1., 0.5], [0., np.nan]],
+                [[1., 0.5], [0., np.nan]],
+                [[0.5, 0.5], [0., 0.5]]]),
         }
         cls.foreground_data = foreground_data
 
     @mock.patch("satpy.composites.enhance2dataset", _enhance2dataset)
     @pytest.mark.parametrize(
         ("foreground_bands", "background_bands", "exp_bands", "exp_result"),
         [
-            ("L", "L", "L", np.array([[1.0, 0.5], [0.0, 1.0]])),
-            ("LA", "LA", "L", np.array([[1.0, 0.75], [0.5, 1.0]])),
-            ("RGB", "RGB", "RGB", np.array([
+            ("L", "L", "L", np.array([[1., 0.5], [0., 1.]])),
+            ("L", "RGB", "RGB", np.array([
                 [[1., 0.5], [0., 1.]],
                 [[1., 0.5], [0., 1.]],
                 [[1., 0.5], [0., 1.]]])),
-            ("RGBA", "RGBA", "RGB", np.array([
+            ("LA", "LA", "LA", np.array([
                 [[1., 0.75], [0.5, 1.]],
-                [[1., 0.75], [0.5, 1.]],
-                [[1., 0.75], [0.5, 1.]]])),
-            ("RGBA", "RGB", "RGB", np.array([
+                [[1., 1.], [1., 1.]]])),
+            ("LA", "RGB", "RGB", np.array([
                 [[1., 0.75], [0.5, 1.]],
                 [[1., 0.75], [0.5, 1.]],
                 [[1., 0.75], [0.5, 1.]]])),
+            ("RGB", "RGB", "RGB", np.array([
+                [[1., 0.5], [0., 1.]],
+                [[1., 0.5], [0., 1.]],
+                [[1., 0.5], [0., 1.]]])),
+            ("RGB", "LA", "RGBA", np.array([
+                [[1., 0.5], [0., 1.]],
+                [[1., 0.5], [0., 1.]],
+                [[1., 0.5], [0., 1.]],
+                [[1., 1.], [1., 1.]]])),
+            ("RGB", "RGBA", "RGBA", np.array([
+                [[1., 0.5], [0., 1.]],
+                [[1., 0.5], [0., 1.]],
+                [[1., 0.5], [0., 1.]],
+                [[1., 1.], [1., 1.]]])),
+            ("RGBA", "RGBA", "RGBA", np.array([
+                [[1., 0.75], [1., 1.]],
+                [[1., 0.75], [1., 1.]],
+                [[1., 0.75], [1., 1.]],
+                [[1., 1.], [1., 1.]]])),
+            ("RGBA", "RGB", "RGB", np.array([
+                [[1., 0.75], [1., 1.]],
+                [[1., 0.75], [1., 1.]],
+                [[1., 0.75], [1., 1.]]])),
         ]
     )
     def test_call(self, foreground_bands, background_bands, exp_bands, exp_result):
@@ -1518,6 +1539,7 @@ def test_call(self, foreground_bands, background_bands, exp_bands, exp_result):
 
         # L mode images
         foreground_data = self.foreground_data[foreground_bands]
+
         attrs = {"mode": foreground_bands, "area": "foo"}
         foreground = xr.DataArray(da.from_array(foreground_data),
                                   dims=("bands", "y", "x"),
@@ -1527,7 +1549,9 @@ def test_call(self, foreground_bands, background_bands, exp_bands, exp_result):
         background = xr.DataArray(da.ones((len(background_bands), 2, 2)), dims=("bands", "y", "x"),
                                   coords={"bands": [c for c in attrs["mode"]]},
                                   attrs=attrs)
+
         res = comp([foreground, background])
+
         assert res.attrs["area"] == "foo"
         np.testing.assert_allclose(res, exp_result)
         assert res.attrs["mode"] == exp_bands