Add compositor for adding an image as a background

.gitignore

@@ -11,6 +11,7 @@ dist
 build
 doc/build
 eggs
+*.eggs
 parts
 bin
 var
@@ -27,6 +28,7 @@ pip-log.txt
 .coverage
 .tox
 nosetests.xml
+htmlcov
 
 #Translations
 *.mo

satpy/composites/__init__.py

@@ -994,9 +994,19 @@ def add_bands(data, bands):
     # Add R, G and B bands, remove L band
     if 'L' in data['bands'].data and 'R' in bands.data:
         lum = data.sel(bands='L')
-        new_data = xr.concat((lum, lum, lum), dim='bands')
-        new_data['bands'] = ['R', 'G', 'B']
-        data = new_data
+        # Keep 'A' if it was present
+        if 'A' in data['bands']:
+            alpha = data.sel(bands='A')
+            new_data = (lum, lum, lum, alpha)
+            new_bands = ['R', 'G', 'B', 'A']
+            mode = 'RGBA'
+        else:
+            new_data = (lum, lum, lum)
+            new_bands = ['R', 'G', 'B']
+            mode = 'RGB'
+        data = xr.concat(new_data, dim='bands', coords={'bands': new_bands})
+        data['bands'] = new_bands
+        data.attrs['mode'] = mode
     # Add alpha band
     if 'A' not in data['bands'].data and 'A' in bands.data:
         new_data = [data.sel(bands=band) for band in data['bands'].data]
@@ -1009,6 +1019,7 @@ def add_bands(data, bands):
         alpha['bands'] = 'A'
         new_data.append(alpha)
         new_data = xr.concat(new_data, dim='bands')
+        new_data.attrs['mode'] = data.attrs['mode'] + 'A'
         data = new_data
 
     return data
@@ -1393,3 +1404,94 @@ def __call__(self, projectables, *args, **kwargs):
         rgb_img = enhance2dataset(projectables[1])
         rgb_img *= luminance
         return super(SandwichCompositor, self).__call__(rgb_img, *args, **kwargs)
+
+
+class StaticImageCompositor(GenericCompositor):
+    """A compositor that loads a static image from disk."""
+
+    def __init__(self, name, fname=None, area=None, **kwargs):
+        """Collect custom configuration values.
+
+        Args:
+            fname (str): Filename of the image to load
+            area (str): Name of area definition for the image.  Optional
+                        for images with built-in area definitions (geotiff)
+        """
+        if fname is None:
+            raise ValueError("No image configured for static image compositor")
+        self.fname = fname
+        self.area = None
+        if area is not None:
+            from satpy.resample import get_area_def
+            self.area = get_area_def(area)
+
+        super(StaticImageCompositor, self).__init__(name, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        from satpy import Scene
+        scn = Scene(reader='generic_image', filenames=[self.fname])
+        scn.load(['image'])
+        img = scn['image']
+        # use compositor parameters as extra metadata
+        # most important: set 'name' of the image
+        img.attrs.update(self.attrs)
+        # Check for proper area definition.  Non-georeferenced images
+        # will raise IndexError
+        try:
+            _ = img.area.size
+        except IndexError:
+            if self.area is None:
+                raise AttributeError("Area definition needs to be configured")
+            img.attrs['area'] = self.area
+        img.attrs['sensor'] = None
+        img.attrs['mode'] = ''.join(img.bands.data)
+        img.attrs.pop('modifiers', None)
+        img.attrs.pop('calibration', None)
+        # Add start time if not present in the filename
+        if 'start_time' not in img.attrs or not img.attrs['start_time']:
+            import datetime as dt
+            img.attrs['start_time'] = dt.datetime.utcnow()
+        if 'end_time' not in img.attrs or not img.attrs['end_time']:
+            import datetime as dt
+            img.attrs['end_time'] = dt.datetime.utcnow()
+
+        return img
+
+
+class BackgroundCompositor(GenericCompositor):
+    """A compositor that overlays one composite on top of another."""
+
+    def __call__(self, projectables, *args, **kwargs):
+        projectables = self.check_areas(projectables)
+
+        # Get enhanced datasets
+        foreground = enhance2dataset(projectables[0])
+        background = enhance2dataset(projectables[1])
+
+        # Adjust bands so that they match
+        # L/RGB -> RGB/RGB
+        # LA/RGB -> RGBA/RGBA
+        # RGB/RGBA -> RGBA/RGBA
+        foreground = add_bands(foreground, background['bands'])
+        background = add_bands(background, foreground['bands'])
+
+        # Get merged metadata
+        attrs = combine_metadata(foreground, background)
+
+        # Stack the images
+        if 'A' in foreground.mode:
+            # Use alpha channel as weight and blend the two composites
+            alpha = foreground.sel(bands='A')
+            data = []
+            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 = xr.where(foreground.isnull(), background, foreground)
+            # Split to separate bands so the mode is correct
+            data = [data.sel(bands=b) for b in data['bands']]
+
+        return super(BackgroundCompositor, self).__call__(data, **kwargs)

satpy/node.py

@@ -26,6 +26,8 @@
 from satpy.utils import get_logger
 
 LOG = get_logger(__name__)
+# Empty leaf used for marking composites with no prerequisites
+EMPTY_LEAF_NAME = "__EMPTY_LEAF_SENTINEL__"
 
 
 class Node(object):
@@ -98,7 +100,9 @@ def display(self, previous=0, include_data=False):
 
     def leaves(self, unique=True):
         """Get the leaves of the tree starting at this root."""
-        if not self.children:
+        if self.name is EMPTY_LEAF_NAME:
+            return []
+        elif not self.children:
             return [self]
         else:
             res = list()
@@ -113,7 +117,7 @@ def trunk(self, unique=True):
         # uniqueness is not correct in `trunk` yet
         unique = False
         res = []
-        if self.children:
+        if self.children and self.name is not EMPTY_LEAF_NAME:
             if self.name is not None:
                 res.append(self)
             for child in self.children:
@@ -154,6 +158,8 @@ def __init__(self, readers, compositors, modifiers):
         # keep a flat dictionary of nodes contained in the tree for better
         # __contains__
         self._all_nodes = DatasetDict()
+        # simplify future logic by only having one "sentinel" empty node
+        self.empty_node = Node(EMPTY_LEAF_NAME)
 
     def leaves(self, nodes=None, unique=True):
         """Get the leaves of the tree starting at this root.
@@ -205,6 +211,9 @@ def add_child(self, parent, child):
         #               but they should all map to the same Node object.
         if self.contains(child.name):
             assert self._all_nodes[child.name] is child
+        if child is self.empty_node:
+            # No need to store "empty" nodes
+            return
         self._all_nodes[child.name] = child
 
     def add_leaf(self, ds_id, parent=None):
@@ -330,6 +339,10 @@ def _get_compositor_prereqs(self, parent, prereq_names, skip=False,
         """
         prereq_ids = []
         unknowns = set()
+        if not prereq_names and not skip:
+            # this composite has no required prerequisites
+            prereq_names = [None]
+
         for prereq in prereq_names:
             n, u = self._find_dependencies(prereq, **dfilter)
             if u:
@@ -418,6 +431,10 @@ def _find_dependencies(self, dataset_key, **dfilter):
                               `satpy.readers.get_key` for more details.
 
         """
+        # Special case: No required dependencies for this composite
+        if dataset_key is None:
+            return self.empty_node, set()
+
         # 0 check if the *exact* dataset is already loaded
         try:
             node = self.getitem(dataset_key)

satpy/scene.py

@@ -752,7 +752,10 @@ def _get_prereq_datasets(self, comp_id, prereq_nodes, keepables, skip=False):
                     and not prereq_node.is_leaf:
                 self._generate_composite(prereq_node, keepables)
 
-            if prereq_id in self.datasets:
+            if prereq_node is self.dep_tree.empty_node:
+                # empty sentinel node - no need to load it
+                continue
+            elif prereq_id in self.datasets:
                 prereq_datasets.append(self.datasets[prereq_id])
             elif not prereq_node.is_leaf and prereq_id in keepables:
                 delayed_gen = True

satpy/tests/compositor_tests/__init__.py

@@ -787,6 +787,60 @@ def test_call(self):
         self.assertEqual(res.attrs['mode'], 'LA')
 
 
+class TestAddBands(unittest.TestCase):
+
+    def test_add_bands(self):
+        from satpy.composites import add_bands
+        import dask.array as da
+        import numpy as np
+        import xarray as xr
+
+        # L + RGB -> RGB
+        data = xr.DataArray(da.ones((1, 3, 3)), dims=('bands', 'y', 'x'),
+                            coords={'bands': ['L']})
+        new_bands = xr.DataArray(da.array(['R', 'G', 'B']), dims=('bands'),
+                                 coords={'bands': ['R', 'G', 'B']})
+        res = add_bands(data, new_bands)
+        res_bands = ['R', 'G', 'B']
+        self.assertEqual(res.mode, ''.join(res_bands))
+        np.testing.assert_array_equal(res.bands, res_bands)
+        np.testing.assert_array_equal(res.coords['bands'], res_bands)
+
+        # L + RGBA -> RGBA
+        data = xr.DataArray(da.ones((1, 3, 3)), dims=('bands', 'y', 'x'),
+                            coords={'bands': ['L']}, attrs={'mode': 'L'})
+        new_bands = xr.DataArray(da.array(['R', 'G', 'B', 'A']), dims=('bands'),
+                                 coords={'bands': ['R', 'G', 'B', 'A']})
+        res = add_bands(data, new_bands)
+        res_bands = ['R', 'G', 'B', 'A']
+        self.assertEqual(res.mode, ''.join(res_bands))
+        np.testing.assert_array_equal(res.bands, res_bands)
+        np.testing.assert_array_equal(res.coords['bands'], res_bands)
+
+        # LA + RGB -> RGBA
+        data = xr.DataArray(da.ones((2, 3, 3)), dims=('bands', 'y', 'x'),
+                            coords={'bands': ['L', 'A']}, attrs={'mode': 'LA'})
+        new_bands = xr.DataArray(da.array(['R', 'G', 'B']), dims=('bands'),
+                                 coords={'bands': ['R', 'G', 'B']})
+        res = add_bands(data, new_bands)
+        res_bands = ['R', 'G', 'B', 'A']
+        self.assertEqual(res.mode, ''.join(res_bands))
+        np.testing.assert_array_equal(res.bands, res_bands)
+        np.testing.assert_array_equal(res.coords['bands'], res_bands)
+
+        # RGB + RGBA -> RGBA
+        data = xr.DataArray(da.ones((3, 3, 3)), dims=('bands', 'y', 'x'),
+                            coords={'bands': ['R', 'G', 'B']},
+                            attrs={'mode': 'RGB'})
+        new_bands = xr.DataArray(da.array(['R', 'G', 'B', 'A']), dims=('bands'),
+                                 coords={'bands': ['R', 'G', 'B', 'A']})
+        res = add_bands(data, new_bands)
+        res_bands = ['R', 'G', 'B', 'A']
+        self.assertEqual(res.mode, ''.join(res_bands))
+        np.testing.assert_array_equal(res.bands, res_bands)
+        np.testing.assert_array_equal(res.coords['bands'], res_bands)
+
+
 def suite():
     """Test suite for all reader tests."""
     loader = unittest.TestLoader()
@@ -809,6 +863,7 @@ def suite():
     mysuite.addTest(loader.loadTestsFromTestCase(TestGenericCompositor))
     mysuite.addTest(loader.loadTestsFromTestCase(TestNIRReflectance))
     mysuite.addTest(loader.loadTestsFromTestCase(TestPrecipCloudsCompositor))
+    mysuite.addTest(loader.loadTestsFromTestCase(TestAddBands))
 
     return mysuite