maps can now plot rgb images with alpha channel (#29)

fmaussion · Nov 2, 2016 · 8b20ed0 · 8b20ed0
1 parent a4f1706
commit 8b20ed0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/salem/graphics.py b/salem/graphics.py
@@ -886,12 +886,12 @@ def set_rgb(self, img=None, crs=None, interp='nearest',
                         pixel_ref='corner').center_grid
             return self.set_rgb(img, grid, interp='linear')
 
-        if (len(img.shape) != 3) or (img.shape[-1] != 3):
-            raise ValueError('img should be of shape (y, x, 3)')
+        if (len(img.shape) != 3) or (img.shape[-1] not in [3, 4]):
+            raise ValueError('img should be of shape (y, x, 3) or (y, x, 4)')
 
         # Unefficient but by far easiest right now
         out = []
-        for i in [0, 1, 2]:
+        for i in range(img.shape[-1]):
             out.append(self._check_data(img[..., i], crs=crs, interp=interp))
         self._rgb = np.dstack(out)