From a36411f7eccd9e68008425298093369a78f231de Mon Sep 17 00:00:00 2001
From: Daniel Wiesmann <daniel.wiesmann@gmail.com>
Date: Wed, 13 Jul 2016 18:52:44 +0100
Subject: [PATCH] Added rgb option to algebra endpoint.

---
 raster/views.py       | 50 ++++++++++++++++++++++++++++++++++++++++---
 tests/test_algebra.py |  4 ++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/raster/views.py b/raster/views.py
index 57d4cd71..acfd9fa9 100644
--- a/raster/views.py
+++ b/raster/views.py
@@ -162,8 +162,6 @@ def get_ids(self):
         return ids
 
     def get(self, request, *args, **kwargs):
-        parser = RasterAlgebraParser()
-
         # Get layer ids
         ids = self.get_ids()
 
@@ -180,7 +178,21 @@ def get(self, request, *args, **kwargs):
                 return self.write_img_to_response(img, {})
 
         # Get formula from request
-        formula = request.GET.get('formula')
+        formula = request.GET.get('formula', None)
+
+        # Dispatch by request type. If a formula was provided, use raster
+        # algebra otherwise look for rgb request.
+        if formula:
+            return self.get_algebra(data, formula)
+        elif 'r' in data and 'g' in data and 'b' in data:
+            return self.get_rgb(data)
+        else:
+            raise RasterAlgebraException(
+                'Specify raster algebra formula or provide rgb layer keys.'
+            )
+
+    def get_algebra(self, data, formula):
+        parser = RasterAlgebraParser()
 
         # Evaluate raster algebra expression, return 400 if not successful
         try:
@@ -216,6 +228,38 @@ def get(self, request, *args, **kwargs):
         # Return rendered image
         return self.write_img_to_response(img, stats)
 
+    def get_rgb(self, data):
+        red = data['r']
+        green = data['g']
+        blue = data['b']
+
+        red = red.bands[0].data()
+        green = green.bands[0].data()
+        blue = blue.bands[0].data()
+
+        scale = float(self.request.GET.get('scale', 255))
+
+        red[red > scale] = scale
+        green[green > scale] = scale
+        blue[blue > scale] = scale
+
+        red = red * 255.0 / scale
+        green = green * 255.0 / scale
+        blue = blue * 255.0 / scale
+
+        # Create zeros array.
+        alpha = 255 * (red > 0) * (blue > 0) * (green > 0)
+
+        rgba = numpy.array((red.ravel(), green.ravel(), blue.ravel(), alpha.ravel())).T
+        rgba = rgba.reshape(WEB_MERCATOR_TILESIZE, WEB_MERCATOR_TILESIZE, 4).astype('uint8')
+
+        # Create image from array
+        img = Image.fromarray(rgba)
+        stats = {}
+
+        # Return rendered image
+        return self.write_img_to_response(img, stats)
+
 
 class TmsView(RasterView):
 
diff --git a/tests/test_algebra.py b/tests/test_algebra.py
index 3e268d2b..95742025 100644
--- a/tests/test_algebra.py
+++ b/tests/test_algebra.py
@@ -120,3 +120,7 @@ def test_band_level_algebra_request(self):
         # Try getting band that does not exist.
         response = self.client.get(self.algebra_tile_url + '?layers=a{0}1={1}&formula=a'.format(BAND_INDEX_SEPARATOR, self.rasterlayer.id))
         self.assertEqual(response.status_code, 400)
+
+    def test_rgb_request(self):
+        response = self.client.get(self.algebra_tile_url + '?layers=r={0},g={0},b={0}'.format(self.rasterlayer.id))
+        self.assertEqual(response.status_code, 200)