From b1ce0fa9e1db3978f6a7f21677869164bc25bbdb Mon Sep 17 00:00:00 2001
From: vincentsarago <vincent.sarago@gmail.com>
Date: Thu, 13 Jun 2024 15:48:47 +0200
Subject: [PATCH] fix alpha band for GCPS and make sure mask are forwarded

---
 CHANGES.md                        |   5 ++
 rio_cogeo/cogeo.py                |  21 +++++-
 tests/fixtures/image_rgb_mask.tif | Bin 790050 -> 790240 bytes
 tests/test_cogeo.py               | 110 +++++++++++++++++++++++++++++-
 4 files changed, 131 insertions(+), 5 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 6c10f370..043f7959 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,10 @@
 # Release Notes
 
+## 5.3.2 (2024-06-13)
+
+* do not set output size in the intermediate VRT
+* add Alpha band for GCPS wrapped dataset
+
 ## 5.3.1 (2024-06-12)
 
 * fix issue when creating COG from file with internal GCPS
diff --git a/rio_cogeo/cogeo.py b/rio_cogeo/cogeo.py
index 85314ae0..84e6b551 100644
--- a/rio_cogeo/cogeo.py
+++ b/rio_cogeo/cogeo.py
@@ -227,6 +227,13 @@ def cog_translate(  # noqa: C901
                     "Cannot add a colormap for multiple bands data."
                 )
 
+            # we need to remove the `alpha band` index from the output data
+            # when we translate the band to an internal mask
+            if alpha and add_mask:
+                indexes = tuple(
+                    x for x in indexes if x in utils.non_alpha_indexes(src_dst)
+                )
+
             if not add_mask and (
                 (nodata is not None or alpha)
                 and dst_kwargs.get("compress", "").lower() == "jpeg"
@@ -241,10 +248,11 @@ def cog_translate(  # noqa: C901
                     else indexes
                 )
 
+            src_indexes = indexes
+
             tilesize = min(int(dst_kwargs["blockxsize"]), int(dst_kwargs["blockysize"]))
 
             vrt_params = {
-                "add_alpha": True,
                 "dtype": dtype,
                 "width": src_dst.width,
                 "height": src_dst.height,
@@ -253,6 +261,7 @@ def cog_translate(  # noqa: C901
             if src_dst.gcps[1]:
                 vrt_params.update(
                     {
+                        "add_alpha": True,
                         "src_crs": src_dst.gcps[1],
                         "src_transform": transform_from_gcps(src_dst.gcps[0]),
                     }
@@ -263,9 +272,12 @@ def cog_translate(  # noqa: C901
                     {"nodata": nodata, "add_alpha": False, "src_nodata": nodata}
                 )
 
-            if alpha:
+            elif alpha:
                 vrt_params.update({"add_alpha": False})
 
+            elif mask:
+                vrt_params.update({"add_alpha": True})
+
             if tms:
                 wo_params = utils.get_web_optimized_params(
                     src_dst,
@@ -276,6 +288,9 @@ def cog_translate(  # noqa: C901
                 )
                 vrt_params.update(**wo_params)
 
+            if vrt_params.get("add_alpha", False) and not (add_mask or mask):
+                indexes = tuple(indexes) + (src_dst.count + 1,)
+
             with WarpedVRT(src_dst, **vrt_params) as vrt_dst:
                 meta = vrt_dst.meta
                 meta["count"] = len(indexes)
@@ -371,7 +386,7 @@ def cog_translate(  # noqa: C901
                 if not quiet:
                     click.echo("Updating dataset tags...", err=True)
 
-                for i, b in enumerate(indexes):
+                for i, b in enumerate(src_indexes):
                     tmp_dst.set_band_description(i + 1, src_dst.descriptions[b - 1])
                     if forward_band_tags:
                         tmp_dst.update_tags(bidx=i + 1, **src_dst.tags(b))
diff --git a/tests/fixtures/image_rgb_mask.tif b/tests/fixtures/image_rgb_mask.tif
index 686a3496d3518f5d988e19c869f0156d82d89702..0f5440e68c21fe93251b90174f938f5ec9c86805 100644
GIT binary patch
delta 340
zcmZ29$Kb(SgAE&4RfQQC7zpCc8(E7PwRnI^mH2oVOo7+~h@*ix4~QFqcnT0N2jX2o
ze3q}dh`qgtoe_wcfS4JGS%8=oh}nP`B*y{7oIuP4#N6A9*m*Xv)mO|(PFTS5;s0`v
U-v$@>8CJe$b*HD0J%v^R0LmnI&Hw-a

delta 114
zcmaDb*I>~cgAE&4H$P!bVVwMd)kKSjfq~%#2M@znAZFs^Vc-X186egKVlyCi2I2r9
sj^%7FVs9^EX9Qv<AZ7+)79eH?Vm2TK$#DQNClGT1G57W&cAgDv07GdS!2kdN

diff --git a/tests/test_cogeo.py b/tests/test_cogeo.py
index b1464461..e9bd60c6 100644
--- a/tests/test_cogeo.py
+++ b/tests/test_cogeo.py
@@ -198,6 +198,35 @@ def test_cog_translate_validAlpha(runner):
                 assert src.compression.value == "JPEG"
                 assert has_mask_band(src)
 
+        cog_translate(
+            raster_path_rgba,
+            "cogeo.tif",
+            cog_profiles.get("deflate"),
+            quiet=True,
+        )
+        with rasterio.open("cogeo.tif") as cog, rasterio.open(
+            raster_path_rgba
+        ) as source:
+            assert cog.read(1, masked=True).max() == source.read(1, masked=True).max()
+            assert cog.count == source.count
+            assert cog.colorinterp == source.colorinterp
+
+        cog_translate(
+            raster_path_rgba,
+            "cogeo.tif",
+            cog_profiles.get("deflate"),
+            quiet=True,
+            add_mask=True,
+        )
+        with rasterio.open("cogeo.tif") as cog, rasterio.open(
+            raster_path_rgba
+        ) as source:
+            assert cog.read(1, masked=True).max() == source.read(1, masked=True).max()
+            assert cog.count <= source.count
+            assert cog.colorinterp != source.colorinterp
+            assert has_mask_band(cog)
+            assert not has_alpha_band(cog)
+
 
 def test_cog_translate_valiNodataNan(runner):
     """Should work as expected and create mask from NaN."""
@@ -274,8 +303,19 @@ def test_cog_translate_mask(runner):
     """Should work as expected (copy mask from input)."""
     with runner.isolated_filesystem():
         cog_translate(raster_path_mask, "cogeo.tif", jpeg_profile, quiet=True)
-        with rasterio.open("cogeo.tif") as src:
-            assert has_mask_band(src)
+        with rasterio.open("cogeo.tif") as cog, rasterio.open(
+            raster_path_mask
+        ) as source:
+            assert cog.read(1, masked=True).max() == source.read(1, masked=True).max()
+            assert cog.count == source.count
+            assert cog.colorinterp == source.colorinterp
+            assert has_mask_band(cog)
+
+            arr = cog.read(1, masked=True)
+            cog_mask = arr.mask
+            arr = source.read(1, masked=True)
+            source_mask = arr.mask
+            numpy.testing.assert_array_equal(cog_mask, source_mask)
 
 
 def test_cog_translate_tags(runner):
@@ -577,6 +617,7 @@ def test_temporaryRaster(fname, is_local, runner):
         raster_path_nan,
         raster_path_nodata,
         raster_path_float,
+        raster_path_gcps,
     ],
 )
 def test_gdal_cog(src_path, runner):
@@ -768,14 +809,79 @@ def test_cog_translate_gcps(runner):
             quiet=True,
         )
 
+        with rasterio.open("cogeo.tif") as cog, rasterio.open(
+            raster_path_gcps
+        ) as source:
+            assert cog.read(1).max() == source.read(1).max()
+            assert not cog.count == source.count
+
+            assert source.gcps[1] is not None
+            # TODO: when we use rio-cogeo, we're using WarpedVRT for the intermediate
+            # step. This result on the output COG to be `reprojected` automatically
+            # ref: https://github.com/cogeotiff/rio-cogeo/issues/292
+            assert cog.gcps[1] is None
+            # we add an alpha band
+            assert cog.count == 2
+            assert cog.colorinterp == (ColorInterp.gray, ColorInterp.alpha)
+
+        cog_translate(
+            raster_path_gcps,
+            "cogeo.tif",
+            cog_profiles.get("deflate"),
+            add_mask=True,
+            quiet=True,
+        )
+
         with rasterio.open("cogeo.tif") as cog, rasterio.open(
             raster_path_gcps
         ) as source:
             assert cog.read(1).max() == source.read(1).max()
             assert cog.count == source.count
+            assert cog.count == 1
+            assert cog.colorinterp == (ColorInterp.gray,)
+            assert has_mask_band(cog)
 
             assert source.gcps[1] is not None
             # TODO: when we use rio-cogeo, we're using WarpedVRT for the intermediate
             # step. This result on the output COG to be `reprojected` automatically
             # ref: https://github.com/cogeotiff/rio-cogeo/issues/292
             assert cog.gcps[1] is None
+
+
+@pytest.mark.parametrize(
+    "src_path",
+    [
+        raster_path_rgb,
+        raster_path_nodata,
+        raster_path_missingnodata,
+        raster_path_mask,
+        raster_path_small,
+    ],
+)
+def test_cog_values(src_path, runner):
+    """Test that COG values are the same."""
+    with runner.isolated_filesystem():
+        cog_translate(
+            src_path,
+            "cogeo.tif",
+            cog_profiles.get("deflate"),
+            quiet=True,
+        )
+        with rasterio.open("cogeo.tif") as cog, rasterio.open(src_path) as source:
+            assert cog.read(1, masked=True).max() == source.read(1, masked=True).max()
+            assert cog.count == source.count
+            assert cog.colorinterp == source.colorinterp
+
+        cog_translate(
+            src_path,
+            "cogeo.tif",
+            cog_profiles.get("deflate"),
+            quiet=True,
+            add_mask=True,
+        )
+        with rasterio.open("cogeo.tif") as cog, rasterio.open(src_path) as source:
+            assert cog.read(1, masked=True).max() == source.read(1, masked=True).max()
+            assert cog.count == source.count
+            assert cog.colorinterp == source.colorinterp
+            assert has_mask_band(cog)
+            assert not has_alpha_band(cog)