Add rgb and rgba webp and manually allow defining colorinterp

kylebarron · kylebarron · commit dcfc0e91a2e2 · 2026-01-29T17:07:12.000-05:00
diff --git a/rasterio_generated/fixtures/uint8_rgb_deflate_block64_cog.py b/rasterio_generated/fixtures/uint8_rgb_deflate_block64_cog.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import numpy as np
+from rasterio.enums import ColorInterp
 
 from rasterio_generated.write_utils import write_cog
 
@@ -27,4 +28,5 @@ def generate(output_path: Path) -> None:
         data,
         blocksize=64,
         compress="DEFLATE",
+        colorinterp=[ColorInterp.red, ColorInterp.green, ColorInterp.blue],
     )
diff --git a/rasterio_generated/fixtures/uint8_rgb_webp_block64_cog.py b/rasterio_generated/fixtures/uint8_rgb_webp_block64_cog.py
@@ -0,0 +1,32 @@
+"""Generate an RGB GeoTIFF with DEFLATE compression."""
+
+from pathlib import Path
+
+import numpy as np
+from rasterio.enums import ColorInterp
+
+from rasterio_generated.write_utils import write_cog
+
+
+def generate(output_path: Path) -> None:
+    """Generate a 256x256 RGB GeoTIFF with DEFLATE compression."""
+    # Create RGB gradient pattern
+    r = np.linspace(0, 127, 128, dtype=np.uint8)
+    g = np.linspace(127, 0, 128, dtype=np.uint8)
+    b = np.full(128, 128, dtype=np.uint8)
+
+    data = np.stack(
+        [
+            np.tile(r, (128, 1)),
+            np.tile(g.reshape(-1, 1), (1, 128)),
+            np.tile(b, (128, 1)),
+        ]
+    )
+
+    write_cog(
+        output_path,
+        data,
+        blocksize=64,
+        compress="WEBP",
+        colorinterp=[ColorInterp.red, ColorInterp.green, ColorInterp.blue],
+    )
diff --git a/rasterio_generated/fixtures/uint8_rgb_webp_block64_cog.tif b/rasterio_generated/fixtures/uint8_rgb_webp_block64_cog.tif
diff --git a/rasterio_generated/fixtures/uint8_rgb_webp_block64_cog_info.md b/rasterio_generated/fixtures/uint8_rgb_webp_block64_cog_info.md
@@ -0,0 +1,55 @@
+```
+Driver: GTiff
+File: rasterio_generated/fixtures/uint8_rgb_webp_block64_cog.tif
+COG: True
+Compression: WEBP
+ColorSpace: None
+
+Profile
+    Width:            128
+    Height:           128
+    Bands:            3
+    Tiled:            True
+    Dtype:            uint8
+    NoData:           0.0
+    Alpha Band:       False
+    Internal Mask:    False
+    Interleave:       PIXEL
+    ColorMap:         False
+    ColorInterp:      ('red', 'green', 'blue')
+    Scales:           (1.0, 1.0, 1.0)
+    Offsets:          (0.0, 0.0, 0.0)
+
+Geo
+    Crs:              EPSG:4326
+    Origin:           (0.0, 0.0)
+    Resolution:       (0.01, -0.01)
+    BoundingBox:      (0.0, -1.28, 1.28, 0.0)
+    MinZoom:          7
+    MaxZoom:          7
+
+Image Metadata
+    AREA_OR_POINT: Area
+
+Image Structure
+    LAYOUT: COG
+    COMPRESSION: WEBP
+    INTERLEAVE: PIXEL
+    OVERVIEW_RESAMPLING: BILINEAR
+    COMPRESSION_REVERSIBILITY: LOSSY
+    WEBP_LEVEL: 75
+
+Band 1
+    ColorInterp: red
+
+Band 2
+    ColorInterp: green
+
+Band 3
+    ColorInterp: blue
+
+IFD
+    Id      Size           BlockSize     Decimation
+    0       128x128        64x64         0
+    1       64x64          64x64         2
+```
diff --git a/rasterio_generated/fixtures/uint8_rgba_webp_block64_cog.py b/rasterio_generated/fixtures/uint8_rgba_webp_block64_cog.py
@@ -0,0 +1,39 @@
+"""Generate an RGB GeoTIFF with DEFLATE compression."""
+
+from pathlib import Path
+
+import numpy as np
+from rasterio.enums import ColorInterp
+
+from rasterio_generated.write_utils import write_cog
+
+
+def generate(output_path: Path) -> None:
+    """Generate a 256x256 RGB GeoTIFF with DEFLATE compression."""
+    # Create RGB gradient pattern
+    r = np.linspace(0, 127, 128, dtype=np.uint8)
+    g = np.linspace(127, 0, 128, dtype=np.uint8)
+    b = np.full(128, 128, dtype=np.uint8)
+    a = np.full(128, 255, dtype=np.uint8)
+
+    data = np.stack(
+        [
+            np.tile(r, (128, 1)),
+            np.tile(g.reshape(-1, 1), (1, 128)),
+            np.tile(b, (128, 1)),
+            np.tile(a.reshape(-1, 1), (1, 128)),
+        ]
+    )
+
+    write_cog(
+        output_path,
+        data,
+        blocksize=64,
+        compress="WEBP",
+        colorinterp=[
+            ColorInterp.red,
+            ColorInterp.green,
+            ColorInterp.blue,
+            ColorInterp.alpha,
+        ],
+    )
diff --git a/rasterio_generated/fixtures/uint8_rgba_webp_block64_cog.tif b/rasterio_generated/fixtures/uint8_rgba_webp_block64_cog.tif
diff --git a/rasterio_generated/fixtures/uint8_rgba_webp_block64_cog_info.md b/rasterio_generated/fixtures/uint8_rgba_webp_block64_cog_info.md
@@ -0,0 +1,58 @@
+```
+Driver: GTiff
+File: rasterio_generated/fixtures/uint8_rgba_webp_block64_cog.tif
+COG: True
+Compression: WEBP
+ColorSpace: None
+
+Profile
+    Width:            128
+    Height:           128
+    Bands:            4
+    Tiled:            True
+    Dtype:            uint8
+    NoData:           0.0
+    Alpha Band:       True
+    Internal Mask:    False
+    Interleave:       PIXEL
+    ColorMap:         False
+    ColorInterp:      ('red', 'green', 'blue', 'alpha')
+    Scales:           (1.0, 1.0, 1.0, 1.0)
+    Offsets:          (0.0, 0.0, 0.0, 0.0)
+
+Geo
+    Crs:              EPSG:4326
+    Origin:           (0.0, 0.0)
+    Resolution:       (0.01, -0.01)
+    BoundingBox:      (0.0, -1.28, 1.28, 0.0)
+    MinZoom:          7
+    MaxZoom:          7
+
+Image Metadata
+    AREA_OR_POINT: Area
+
+Image Structure
+    LAYOUT: COG
+    COMPRESSION: WEBP
+    INTERLEAVE: PIXEL
+    OVERVIEW_RESAMPLING: BILINEAR
+    COMPRESSION_REVERSIBILITY: LOSSY
+    WEBP_LEVEL: 75
+
+Band 1
+    ColorInterp: red
+
+Band 2
+    ColorInterp: green
+
+Band 3
+    ColorInterp: blue
+
+Band 4
+    ColorInterp: alpha
+
+IFD
+    Id      Size           BlockSize     Decimation
+    0       128x128        64x64         0
+    1       64x64          64x64         2
+```
diff --git a/rasterio_generated/write_utils.py b/rasterio_generated/write_utils.py
@@ -41,6 +41,7 @@ def write_cog(
     predictor: Literal[2, 3] | None = None,
     nodata: int | float | None = None,
     rasterio_env: dict[str, str | bool | int] | None = None,
+    colorinterp: list[ColorInterp] | None = None,
 ):
     """Write a COG to disk.
 
@@ -106,19 +107,17 @@ def write_cog(
     with rasterio.Env(env):
         with MemoryFile() as memfile:
             with memfile.open(**src_profile) as mem:
-                if nband == 3:
-                    ci = [ColorInterp.red, ColorInterp.green, ColorInterp.blue]
-
-                else:
-                    ci = [ColorInterp.gray]
-                    if nband > 1:
-                        ci += [ColorInterp.undefined] * (nband - 1)
+                if colorinterp is not None:
+                    assert len(colorinterp) == mem.count
 
                 if nodata_type == "alpha" and mask is not None:
                     data = np.concatenate([data, mask])
-                    ci += [ColorInterp.alpha]
+                    assert (
+                        colorinterp is not None and colorinterp[-1] == ColorInterp.alpha
+                    )
 
-                mem.colorinterp = ci
+                if colorinterp is not None:
+                    mem.colorinterp = colorinterp
 
                 # Write Data
                 if nband == 1:
