diff --git a/CHANGES.md b/CHANGES.md index 6c10f37..043f795 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 85314ae..84e6b55 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 686a349..0f5440e 100644 Binary files a/tests/fixtures/image_rgb_mask.tif and b/tests/fixtures/image_rgb_mask.tif differ diff --git a/tests/test_cogeo.py b/tests/test_cogeo.py index b146446..e9bd60c 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)