Add color interpolation to crop output, simplify logic

OpenDataAnalytics · Apr 8, 2019 · 6c3b694 · 6c3b694
1 parent 3d8628f
commit 6c3b694
Showing 1 changed file with 12 additions and 17 deletions.
diff --git a/gaia/geo/gdal_functions.py b/gaia/geo/gdal_functions.py
@@ -281,7 +281,8 @@ def world_to_pixel(geoMatrix, x, y):
         return None
 
     if raster_input.RasterCount == 1:
-        clip = src_array[ul_y:lr_y, ul_x:lr_x]
+        clip2d = src_array[ul_y:lr_y, ul_x:lr_x]
+        clip = np.expand_dims(clip2d, axis=0)
     else:
         clip = src_array[:, ul_y:lr_y, ul_x:lr_x]
 
@@ -314,13 +315,9 @@ def world_to_pixel(geoMatrix, x, y):
     mask = image_to_array(raster_poly)
 
     # Clip the image using the mask
-    if raster_input.RasterCount == 1:
-        clip = gdalnumeric.numpy.choose(
-            mask, (clip, nodata_value)).astype(src_dtype)
-    else:
-        for i in range(raster_input.RasterCount):
-            clip[i] = gdalnumeric.numpy.choose(
-                mask, (clip[i], nodata_value)).astype(src_dtype)
+    for i in range(raster_input.RasterCount):
+        clip[i] = gdalnumeric.numpy.choose(
+            mask, (clip[i], nodata_value)).astype(src_dtype)
 
     # create output raster
     raster_band = raster_input.GetRasterBand(1)
@@ -333,15 +330,13 @@ def world_to_pixel(geoMatrix, x, y):
     gdalnumeric.CopyDatasetInfo(raster_input, output_dataset,
                                 xoff=xoffset, yoff=yoffset)
     bands = raster_input.RasterCount
-    if bands > 1:
-        for i in range(bands):
-            outBand = output_dataset.GetRasterBand(i + 1)
-            outBand.SetNoDataValue(nodata_values[i])
-            outBand.WriteArray(clip[i])
-    else:
-        outBand = output_dataset.GetRasterBand(1)
-        outBand.SetNoDataValue(nodata_values[0])
-        outBand.WriteArray(clip)
+    for i in range(bands):
+        inband = raster_input.GetRasterBand(i + 1)
+
+        outBand = output_dataset.GetRasterBand(i + 1)
+        outBand.SetColorInterpretation(inband.GetColorInterpretation())
+        outBand.SetNoDataValue(nodata_values[i])
+        outBand.WriteArray(clip[i])
 
     if raster_output:
         output_driver = gdal.GetDriverByName('GTiff')