refactor: split create_raster_map_layer function and add comments

uvdat/core/tasks/map_layers.py

@@ -46,11 +46,63 @@ def add_styling(geojson_data, style_options):
     return geopandas.GeoDataFrame.from_features(features)
 
 
+def rasterio_convert_raster(
+    map_layer, temp_dir, raw_data_path, transparency_threshold, trim_distribution_percentage
+):
+    raster_path = Path(temp_dir, 'raster.tiff')
+    with open(raw_data_path, 'rb') as raw_data:
+        # read original data with rasterio
+        input_data = rasterio.open(raw_data)
+        output_data = rasterio.open(
+            raster_path,
+            'w',
+            driver='GTiff',
+            height=input_data.height,
+            width=input_data.width,
+            count=1,
+            dtype=numpy.float32,
+            crs=input_data.crs,
+            transform=input_data.transform,
+        )
+        band = input_data.read(1)
+
+        # alter data according to style options
+        if trim_distribution_percentage:
+            # trim a number of values from both ends of the distribution
+            histogram, bin_edges = numpy.histogram(band, bins=1000)
+            trim_n = band.size * trim_distribution_percentage
+            new_min = None
+            new_max = None
+            sum_values = 0
+            for bin_index, bin_count in enumerate(histogram):
+                bin_edge = bin_edges[bin_index]
+                sum_values += bin_count
+                if new_min is None and sum_values > trim_n:
+                    new_min = bin_edge
+                if new_max is None and sum_values > band.size - trim_n:
+                    new_max = bin_edge
+            if new_min:
+                band[band < new_min] = new_min
+            if new_max:
+                band[band > new_max] = new_max
+        if transparency_threshold is not None:
+            # clamp values below transparency threshold
+            band[band < transparency_threshold] = transparency_threshold
+
+        band_range = [float(band.min()), float(band.max())]
+        map_layer.default_style['data_range'] = band_range
+
+        output_data.write(band, 1)
+        output_data.close()
+    return raster_path
+
+
 def create_raster_map_layer(file_item, style_options):
     """Save a RasterMapLayer from a FileItem's contents."""
     import large_image
     import large_image_converter
 
+    # create new raster map layer object
     new_map_layer = RasterMapLayer.objects.create(
         dataset=file_item.dataset,
         metadata={},
@@ -59,81 +111,57 @@ def create_raster_map_layer(file_item, style_options):
     )
     print('\t', f'RasterMapLayer {new_map_layer.id} created.')
 
+    transparency_threshold = style_options.get('transparency_threshold')
+    trim_distribution_percentage = style_options.get('trim_distribution_percentage')
+
     with tempfile.TemporaryDirectory() as temp_dir:
         raster_path = None
         raw_data_path = Path(temp_dir, 'raw_data.tiff')
         cog_raster_path = Path(temp_dir, 'cog_raster.tiff')
+        # write original data from file field to file in temp_dir
         with open(raw_data_path, 'wb') as raw_data:
             with file_item.file.open('rb') as raw_data_archive:
                 raw_data.write(raw_data_archive.read())
 
-        try:
-            source = large_image.open(raw_data_path)
-            if source.geospatial:
-                raster_path = raw_data_path
-                min_val, max_val = None, None
-                source._scanForMinMax(numpy.int64)
-                for _frame, range_spec in source._bandRanges.items():
-                    frame_min = numpy.min(range_spec.get('min', numpy.empty(1)))
-                    frame_max = numpy.max(range_spec.get('max', numpy.empty(1)))
-                    if min_val is None or frame_min < min_val:
-                        min_val = frame_min
-                    if max_val is None or frame_max < max_val:
-                        max_val = frame_max
-                new_map_layer.default_style['data_range'] = [int(min_val), int(max_val)]
-        except large_image.exceptions.TileSourceError:
-            pass
-
-        if raster_path is None:
-            transparency_threshold = style_options.get('transparency_threshold')
-            trim_distribution_percentage = style_options.get('trim_distribution_percentage')
-
-            raster_path = Path(temp_dir, 'raster.tiff')
-            with open(raw_data_path, 'rb') as raw_data:
-                input_data = rasterio.open(raw_data)
-                output_data = rasterio.open(
-                    raster_path,
-                    'w',
-                    driver='GTiff',
-                    height=input_data.height,
-                    width=input_data.width,
-                    count=1,
-                    dtype=numpy.float32,
-                    crs=input_data.crs,
-                    transform=input_data.transform,
+        if transparency_threshold or trim_distribution_percentage:
+            # if data needs to be altered according to style options, use rasterio
+            raster_path = rasterio_convert_raster(
+                new_map_layer,
+                temp_dir,
+                raw_data_path,
+                transparency_threshold,
+                trim_distribution_percentage,
+            )
+        else:
+            try:
+                # if applicable, use large_image to interpret data
+                source = large_image.open(raw_data_path)
+                if source.geospatial:
+                    raster_path = raw_data_path
+                    min_val, max_val = None, None
+                    source._scanForMinMax(numpy.int64)
+                    for _frame, range_spec in source._bandRanges.items():
+                        frame_min = numpy.min(range_spec.get('min', numpy.empty(1)))
+                        frame_max = numpy.max(range_spec.get('max', numpy.empty(1)))
+                        if min_val is None or frame_min < min_val:
+                            min_val = frame_min
+                        if max_val is None or frame_max < max_val:
+                            max_val = frame_max
+                    new_map_layer.default_style['data_range'] = [int(min_val), int(max_val)]
+            except large_image.exceptions.TileSourceError:
+                # if original data cannot be interpreted by large_image, use rasterio
+                raster_path = rasterio_convert_raster(
+                    new_map_layer,
+                    temp_dir,
+                    raw_data_path,
+                    transparency_threshold,
+                    trim_distribution_percentage,
                 )
-                band = input_data.read(1)
-
-                if trim_distribution_percentage:
-                    # trim a number of values from both ends of the distribution
-                    histogram, bin_edges = numpy.histogram(band, bins=1000)
-                    trim_n = band.size * trim_distribution_percentage
-                    new_min = None
-                    new_max = None
-                    sum_values = 0
-                    for bin_index, bin_count in enumerate(histogram):
-                        bin_edge = bin_edges[bin_index]
-                        sum_values += bin_count
-                        if new_min is None and sum_values > trim_n:
-                            new_min = bin_edge
-                        if new_max is None and sum_values > band.size - trim_n:
-                            new_max = bin_edge
-                    if new_min:
-                        band[band < new_min] = new_min
-                    if new_max:
-                        band[band > new_max] = new_max
-
-                if transparency_threshold is not None:
-                    band[band < transparency_threshold] = transparency_threshold
-
-                band_range = [float(band.min()), float(band.max())]
-                new_map_layer.default_style['data_range'] = band_range
-
-                output_data.write(band, 1)
-                output_data.close()
 
+        # use large_image to convert new raster data to COG
         large_image_converter.convert(str(raster_path), str(cog_raster_path))
         with open(cog_raster_path, 'rb') as cog_raster_file:
+            # save COG to new raster map layer
             new_map_layer.cloud_optimized_geotiff.save(
                 cog_raster_path, ContentFile(cog_raster_file.read())
             )