From 641d5a2acfd30c9a253d8d049207e58a26d5021a Mon Sep 17 00:00:00 2001
From: Vincent Sarago <vincent.sarago@gmail.com>
Date: Thu, 29 Jul 2021 19:13:50 +0200
Subject: [PATCH] Fix alpha band values when storing `Uint16` data in PNG
 (#407)

* Fix alpha band values when storing `Uint16` data in **PNG**

* make sure we keep even distribution of alpha values
---
 CHANGES.md           |  4 +++-
 rio_tiler/utils.py   |  3 +++
 tests/test_models.py | 40 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index 3553ab33..987472da 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,4 +1,4 @@
-## Unreleased
+## 2.1.1 (2021-07-29)
 
 * add support for setting the S3 endpoint url via the `AWS_S3_ENDPOINT` environment variables in `aws_get_object` function using boto3 (https://github.com/cogeotiff/rio-tiler/pull/394)
 * make `ImageStatistics.valid_percent` a value between 0 and 100 (instead of 0 and 1) (author @param-thakker, https://github.com/cogeotiff/rio-tiler/pull/400)
@@ -9,6 +9,8 @@
         pass
     ```
 
+* Fix alpha band values when storing `Uint16` data in **PNG**. (https://github.com/cogeotiff/rio-tiler/pull/407)
+
 ## 2.1.0 (2021-05-17)
 
 * add auto-rescaling in `ImageData.render` method to avoid error when datatype is not supported by the output driver (https://github.com/cogeotiff/rio-tiler/pull/391)
diff --git a/rio_tiler/utils.py b/rio_tiler/utils.py
index c6369ad0..c0461ca8 100644
--- a/rio_tiler/utils.py
+++ b/rio_tiler/utils.py
@@ -344,6 +344,9 @@ def render(
     if img_format == "WEBP" and data.shape[0] == 1:
         data = numpy.repeat(data, 3, axis=0)
 
+    if img_format == "PNG" and data.dtype == "uint16" and mask is not None:
+        mask = linear_rescale(mask, (0, 255), (0, 65535)).astype("uint16")
+
     elif img_format == "JPEG":
         mask = None
 
diff --git a/tests/test_models.py b/tests/test_models.py
index f4d1c2bb..1491a2e1 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,7 +1,10 @@
 """Test rio_tiler.models."""
 
+from io import BytesIO
+
 import numpy
 import pytest
+import rasterio
 
 from rio_tiler.errors import InvalidDatatypeWarning
 from rio_tiler.models import ImageData
@@ -53,3 +56,40 @@ def test_imageData_AutoRescalingAllTypes(dtype):
         ImageData(numpy.zeros((3, 256, 256), dtype=dtype)).render(
             img_format="JP2OPENJPEG"
         )
+
+
+def test_16bit_PNG():
+    """Uint16 Mask value should be between 0 and 65535 for PNG."""
+    mask = numpy.zeros((256, 256), dtype="uint16") + 255
+    mask[0:10, 0:10] = 0
+    mask[10:11, 10:11] = 100
+
+    with pytest.warns(None):
+        img = ImageData(numpy.zeros((1, 256, 256), dtype="uint16"), mask,).render(
+            img_format="PNG"
+        )
+
+        with rasterio.open(BytesIO(img)) as src:
+            assert src.count == 2
+            assert src.meta["dtype"] == "uint16"
+            arr = src.read(2)
+            assert arr.min() == 0
+            assert arr.max() == 65535
+            assert (arr[0:10, 0:10] == 0).all()
+            assert (arr[10:11, 10:11] == 25700).all()
+            assert (arr[11:, 11:] == 65535).all()
+
+    with pytest.warns(None):
+        img = ImageData(numpy.zeros((3, 256, 256), dtype="uint16"), mask,).render(
+            img_format="PNG"
+        )
+
+        with rasterio.open(BytesIO(img)) as src:
+            assert src.count == 4
+            assert src.meta["dtype"] == "uint16"
+            arr = src.read(4)
+            assert arr.min() == 0
+            assert arr.max() == 65535
+            assert (arr[0:10, 0:10] == 0).all()
+            assert (arr[10:11, 10:11] == 25700).all()
+            assert (arr[11:, 11:] == 65535).all()