Reprojection working

2021-11-01 15:51:40 -04:00 · 2021-11-01 15:51:40 -04:00 · 9bcedc65b7
commit 9bcedc65b7
--- a/app/api/tasks.py
+++ b/app/api/tasks.py
@ -360,8 +360,6 @@ class TaskDownloads(TaskNestedView):
        if not os.path.exists(asset_path):
            raise exceptions.NotFound(_("Asset does not exist"))
            
-        print(request.GET.get('filename'))
-
        return download_file_response(request, asset_path, 'attachment', download_filename=request.GET.get('filename'))

 """
--- a/app/api/tiler.py
+++ b/app/api/tiler.py
@ -20,14 +20,14 @@ from rio_tiler.io import COGReader
 from rio_tiler.errors import InvalidColorMapName
 import numpy as np
 from .custom_colormaps_helper import custom_colormaps
-from app.raster_utils import export_raster_index
+from app.raster_utils import export_raster, extension_for_export_format
 from .hsvblend import hsv_blend
 from .hillshade import LightSource
 from .formulas import lookup_formula, get_algorithm_list
 from .tasks import TaskNestedView
 from rest_framework import exceptions
 from rest_framework.response import Response
-from worker.tasks import export_raster_index
+from worker.tasks import export_raster
 from django.utils.translation import gettext as _

 ZOOM_EXTRA_LEVELS = 3
@ -489,15 +489,9 @@ class Export(TaskNestedView):

        expr = None

-        if not export_format in ['gtiff', 'gtiff-rgb', 'jpg']:
+        if not export_format in ['gtiff', 'gtiff-rgb', 'jpg', 'png']:
            raise exceptions.ValidationError(_("Unsupported format: %(format)") % {'format': export_format})
        
-        extensions = {
-            'gtiff': 'tif',
-            'gtiff-rgb': 'tif',
-            'jpg': 'jpg'
-        }
-
        if epsg is not None:
            try:
                epsg = int(epsg)
@ -510,9 +504,12 @@ class Export(TaskNestedView):
            except ValueError as e:
                raise exceptions.ValidationError(str(e))
        
-        if export_format == 'gtiff-rgb' or export_format == 'jpg':
+        if export_format in ['gtiff-rgb', 'jpg', 'png']:
            if formula is not None and rescale is None:
                rescale = "-1,1"
+        
+        if export_format == 'gtiff':
+            rescale = None

        url = get_raster_path(task, asset_type)

@ -520,14 +517,16 @@ class Export(TaskNestedView):
            raise exceptions.NotFound()
        
        # Strip unsafe chars, append suffix
-        filename = "{}.{}".format(
+        extension = extension_for_export_format(export_format)
+        filename = "{}{}.{}".format(
                        re.sub(r'[^0-9a-zA-Z-_]+', '', task.name.replace(" ", "-").replace("/", "-")) + "-" + asset_type,
-                        extensions[export_format]
+                        "-{}".format(formula) if expr is not None else "",
+                        extension
                    )

        # Shortcut the process if no processing is required
        if export_format == 'gtiff' and (epsg == task.epsg or epsg is None) and expr is None:
            return Response({'url': '/api/projects/{}/tasks/{}/download/{}.tif'.format(task.project.id, task.id, asset_type), 'filename': filename})
        else:
-            celery_task_id = export_raster_index.delay(url, expr).task_id
+            celery_task_id = export_raster.delay(url, epsg=epsg, expression=expr, format=export_format).task_id
            return Response({'celery_task_id': celery_task_id, 'filename': filename})
--- a/app/raster_utils.py
+++ b/app/raster_utils.py
@ -5,45 +5,97 @@ import numpy as np
 import numexpr as ne
 from rasterio.enums import ColorInterp
 from rio_tiler.utils import has_alpha_band
+from rasterio.warp import calculate_default_transform, reproject, Resampling

+def extension_for_export_format(export_format):
+    extensions = {
+        'gtiff': 'tif',
+        'gtiff-rgb': 'tif',
+        'jpg': 'jpg',
+        'png': 'png'
+    }
+    return extensions.get(export_format, 'tif')
+
+def export_raster(input, output, **opts):
+    epsg = opts.get('epsg')
+    expression = opts.get('expression')

-def export_raster_index(input, expression, output):
    with rasterio.open(input) as src:
-        profile = src.profile
-        profile.update(
-            dtype=rasterio.float32,
-            count=1,
-            nodata=-9999
-        )
+        profile = src.meta.copy()

-        bands_names = ["b{}".format(b) for b in tuple(sorted(set(re.findall(r"b(?P<bands>[0-9]{1,2})", expression))))]
-        rgb = expression.split(",")
-        indexes = tuple([int(b.replace("b", "")) for b in bands_names])
+        # Define write band function
+        # Reprojection needed?
+        if src.crs is not None and epsg is not None and src.crs.to_epsg() != epsg:
+            dst_crs = "EPSG:{}".format(epsg)

-        alpha_index = None
-        if has_alpha_band(src):
-            try:
-                alpha_index = src.colorinterp.index(ColorInterp.alpha) + 1
-                indexes += (alpha_index, )
-            except ValueError:
-                pass
+            transform, width, height = calculate_default_transform(
+                src.crs, dst_crs, src.width, src.height, *src.bounds)

-        data = src.read(indexes=indexes, out_dtype=np.float32)
-        arr = dict(zip(bands_names, data))
-        arr = np.array([np.nan_to_num(ne.evaluate(bloc.strip(), local_dict=arr)) for bloc in rgb])
+            profile.update(
+                crs=dst_crs,
+                transform=transform,
+                width=width,
+                height=height
+            )

-        # Set nodata values
-        index_band = arr[0]
-        if alpha_index is not None:
-            # -1 is the last band = alpha
-            index_band[data[-1] == 0] = -9999
+            def write_band(arr, dst, i):
+                return reproject(source=arr, 
+                        destination=rasterio.band(dst, i),
+                        src_transform=src.transform,
+                        src_crs=src.crs,
+                        dst_transform=transform,
+                        dst_crs=dst_crs,
+                        resampling=Resampling.nearest)

-        # Remove infinity values
-        index_band[index_band>1e+30] = -9999
-        index_band[index_band<-1e+30] = -9999
+        else:
+            # No reprojection needed
+            def write_band(arr, dst, i):
+                dst.write(arr, i)
+        
+        # TODO: output format
+        profile.update(driver='GTiff')

-        # Make sure this is float32
-        arr = arr.astype(np.float32)
+        if expression is not None:
+            # Apply band math
+            profile.update(
+                dtype=rasterio.float32,
+                count=1,
+                nodata=-9999
+            )

-        with rasterio.open(output, 'w', **profile) as dst:
-            dst.write(arr)
+            bands_names = ["b{}".format(b) for b in tuple(sorted(set(re.findall(r"b(?P<bands>[0-9]{1,2})", expression))))]
+            rgb = expression.split(",")
+            indexes = tuple([int(b.replace("b", "")) for b in bands_names])
+
+            alpha_index = None
+            if has_alpha_band(src):
+                try:
+                    alpha_index = src.colorinterp.index(ColorInterp.alpha) + 1
+                    indexes += (alpha_index, )
+                except ValueError:
+                    pass
+
+            data = src.read(indexes=indexes, out_dtype=np.float32)
+            arr = dict(zip(bands_names, data))
+            arr = np.array([np.nan_to_num(ne.evaluate(bloc.strip(), local_dict=arr)) for bloc in rgb])
+
+            # Set nodata values
+            index_band = arr[0]
+            if alpha_index is not None:
+                # -1 is the last band = alpha
+                index_band[data[-1] == 0] = -9999
+
+            # Remove infinity values
+            index_band[index_band>1e+30] = -9999
+            index_band[index_band<-1e+30] = -9999
+
+            # Make sure this is float32
+            arr = arr.astype(np.float32)
+
+            with rasterio.open(output, 'w', **profile) as dst:
+                write_band(arr, dst, 1)
+        else:
+            # Copy bands as-is
+            with rasterio.open(output, 'w', **profile) as dst:
+                for i in range(1, src.count + 1):
+                    write_band(src.read(i), dst, i)
--- a/app/static/app/js/components/ExportAssetPanel.jsx
+++ b/app/static/app/js/components/ExportAssetPanel.jsx
@ -9,10 +9,11 @@ import Workers from '../classes/Workers';

 export default class ExportAssetPanel extends React.Component {
  static defaultProps = {
-      exportFormats: ["gtiff", "gtiff-rgb", "jpg"],
+      exportFormats: ["gtiff", "gtiff-rgb", "jpg", "png"],
      asset: "",
      exportParams: {},
-      task: null
+      task: null,
+      dropUp: false
  };
  static propTypes = {
      exportFormats: PropTypes.arrayOf(PropTypes.string),
@ -21,7 +22,8 @@ export default class ExportAssetPanel extends React.Component {
          PropTypes.func,
          PropTypes.object
      ]),
-      task: PropTypes.object.isRequired
+      task: PropTypes.object.isRequired,
+      dropUp: PropTypes.bool
  }

  constructor(props){
@ -39,6 +41,10 @@ export default class ExportAssetPanel extends React.Component {
        'jpg': {
            label: _("JPEG (RGB)"),
            icon: "fas fa-palette"
+        },
+        'png': {
+            label: _("PNG (RGB)"),
+            icon: "fas fa-palette"
        }
    };

@ -160,7 +166,7 @@ export default class ExportAssetPanel extends React.Component {
        <div className="row form-group form-inline">
            <label className="col-sm-3 control-label">{_("Export:")}</label>
            <div className="col-sm-9">
-                <div className="btn-group dropup">
+                <div className={"btn-group " + (this.props.dropUp ?  "dropup" : "")}>
                    <button onClick={this.handleExport(exportFormats[0])}
                        disabled={disabled} type="button" className="btn btn-sm btn-primary btn-export">
                        {exporting ? <i className="fa fa-spin fa-circle-notch"/> : <i className={this.efInfo[exportFormats[0]].icon + " fa-fw"}/>} {this.efInfo[exportFormats[0]].label}
--- a/app/static/app/js/components/LayersControlLayer.jsx
+++ b/app/static/app/js/components/LayersControlLayer.jsx
@ -331,7 +331,10 @@ export default class LayersControlLayer extends React.Component {
                </div>
            </div> : ""}

-            <ExportAssetPanel task={meta.task} asset={this.asset} exportParams={this.getLayerParams} />
+            <ExportAssetPanel task={meta.task} 
+                            asset={this.asset} 
+                            exportParams={this.getLayerParams} 
+                            dropUp />
        </div> : ""}
    </div>);

--- a/worker/tasks.py
+++ b/worker/tasks.py
@ -2,6 +2,7 @@ import os
 import shutil
 import tempfile
 import traceback
+import json

 import time
 from threading import Event, Thread
@ -18,7 +19,7 @@ from nodeodm.models import ProcessingNode
 from webodm import settings
 import worker
 from .celery import app
-from app.raster_utils import export_raster_index as export_raster_index_sync
+from app.raster_utils import export_raster as export_raster_sync, extension_for_export_format
 import redis

 logger = get_task_logger("app.logger")
@ -154,11 +155,11 @@ def execute_grass_script(script, serialized_context = {}, out_key='output'):


@app.task(bind=True)
-def export_raster_index(self, input, expression):
+def export_raster(self, input, **opts):
    try:
-        logger.info("Exporting raster index {} with expression: {}".format(input, expression))
-        tmpfile = tempfile.mktemp('_raster_index.tif', dir=settings.MEDIA_TMP)
-        export_raster_index_sync(input, expression, tmpfile)
+        logger.info("Exporting raster {} with options: {}".format(input, json.dumps(opts)))
+        tmpfile = tempfile.mktemp('_raster.{}'.format(extension_for_export_format(opts.get('format', 'gtiff'))), dir=settings.MEDIA_TMP)
+        export_raster_sync(input, tmpfile, **opts)
        result = {'file': tmpfile}

        if settings.TESTING: