Fixing wrapping of numpy functions in browse rendering

EOxServer · Jan 16, 2023 · 3531780 · 3531780
1 parent 64d457a
commit 3531780
Showing 1 changed file with 50 additions and 32 deletions.
diff --git a/eoxserver/render/browse/functions.py b/eoxserver/render/browse/functions.py
@@ -221,6 +221,7 @@ def percentile(ds, perc, default=0):
 def _has_stats(band):
     return 'STATISTICS_MINIMUM' in band.GetMetadata()
 
+
 def statistics_min(ds, default=0):
     band = ds.GetRasterBand(1)
     if _has_stats(band):
@@ -236,6 +237,7 @@ def statistics_max(ds, default=0):
         return max_
     return default
 
+
 def statistics_mean(ds, default=0):
     band = ds.GetRasterBand(1)
     if _has_stats(band):
@@ -295,42 +297,58 @@ def interpolate(
 
 def wrap_numpy_func(function):
     @wraps(function)
-    def inner(ds, *args, **kwargs):
-        band = ds.GetRasterBand(1)
-        data = band.ReadAsArray()
-        function(data, *args, **kwargs)
-        band.WriteArray(data)
-        return ds
+    def inner(*args, **kwargs):
+        converted_args = []
+        for arg in args:
+            if isinstance(arg, gdal.Dataset):
+                band = arg.GetRasterBand(1)
+                data = band.ReadAsArray()
+                converted_args.append(data)
+            else:
+                converted_args.append(arg)
+
+        result = function(*converted_args, **kwargs)
+
+        if np.isscalar(result):
+            return result
+
+        arg = args[0]
+        if isinstance(arg, gdal.Dataset):
+            band = arg.GetRasterBand(1)
+            band.WriteArray(result)
+            return arg
+        else:
+            return gdal_array.OpenNumPyArray(result, False)
     return inner
 
 
 function_map = {
-    'sin': np.sin,
-    'cos': np.cos,
-    'tan': np.tan,
-    'arcsin': np.arcsin,
-    'arccos': np.arccos,
-    'arctan': np.arctan,
-    'hypot': np.hypot,
-    'arctan2': np.arctan2,
-    'degrees': np.degrees,
-    'radians': np.radians,
-    'unwrap': np.unwrap,
-    'deg2rad': np.deg2rad,
-    'rad2deg': np.rad2deg,
-    'sinh': np.sinh,
-    'cosh': np.cosh,
-    'tanh': np.tanh,
-    'arcsinh': np.arcsinh,
-    'arccosh': np.arccosh,
-    'arctanh': np.arctanh,
-    'exp': np.exp,
-    'expm1': np.expm1,
-    'exp2': np.exp2,
-    'log': np.log,
-    'log10': np.log10,
-    'log2': np.log2,
-    'log1p': np.log1p,
+    'sin': wrap_numpy_func(np.sin),
+    'cos': wrap_numpy_func(np.cos),
+    'tan': wrap_numpy_func(np.tan),
+    'arcsin': wrap_numpy_func(np.arcsin),
+    'arccos': wrap_numpy_func(np.arccos),
+    'arctan': wrap_numpy_func(np.arctan),
+    'hypot': wrap_numpy_func(np.hypot),
+    'arctan2': wrap_numpy_func(np.arctan2),
+    'degrees': wrap_numpy_func(np.degrees),
+    'radians': wrap_numpy_func(np.radians),
+    'unwrap': wrap_numpy_func(np.unwrap),
+    'deg2rad': wrap_numpy_func(np.deg2rad),
+    'rad2deg': wrap_numpy_func(np.rad2deg),
+    'sinh': wrap_numpy_func(np.sinh),
+    'cosh': wrap_numpy_func(np.cosh),
+    'tanh': wrap_numpy_func(np.tanh),
+    'arcsinh': wrap_numpy_func(np.arcsinh),
+    'arccosh': wrap_numpy_func(np.arccosh),
+    'arctanh': wrap_numpy_func(np.arctanh),
+    'exp': wrap_numpy_func(np.exp),
+    'expm1': wrap_numpy_func(np.expm1),
+    'exp2': wrap_numpy_func(np.exp2),
+    'log': wrap_numpy_func(np.log),
+    'log10': wrap_numpy_func(np.log10),
+    'log2': wrap_numpy_func(np.log2),
+    'log1p': wrap_numpy_func(np.log1p),
     'hillshade': hillshade,
     'slopeshade': slopeshade,
     'aspect': aspect,