add color interp (#75)

* add color interp

* add palette test case

* parse color map

* handle nodata and alpha band

* color_map -> colormap

* test colormap

aiocogeo/cog.py

@@ -12,12 +12,12 @@
 import numpy as np
 
 from . import config
-from .constants import MaskFlags, PHOTOMETRIC
+from .constants import ColorInterp, MaskFlags, PHOTOMETRIC
 from .errors import InvalidTiffError, TileNotFoundError
 from .filesystems import Filesystem
 from .ifd import IFD, ImageIFD, MaskIFD
 from .partial_reads import PartialReadInterface
-from .utils import run_in_background
+from .utils import run_in_background, chunks
 
 logger = logging.getLogger(__name__)
 logger.setLevel(config.LOG_LEVEL)
@@ -114,7 +114,7 @@ def profile(self) -> Dict[str, Any]:
             "crs": f"EPSG:{self.epsg}",
             "nodata": ifd.nodata,
             "tiled": True,
-            "photometric": PHOTOMETRIC[ifd.PhotometricInterpretation.value],
+            "photometric": self.photometric,
         }
 
     @property
@@ -184,6 +184,57 @@ def mask_flags(self):
             return band_flags
         return [flags for _ in range(bands)]
 
+    @property
+    def photometric(self):
+        return PHOTOMETRIC[self.ifds[0].PhotometricInterpretation.value]
+
+    @property
+    def colormap(self) -> Optional[Dict[int, Tuple[int, int, int]]]:
+        """https://www.awaresystems.be/imaging/tiff/tifftags/colormap.html"""
+        if self.ifds[0].ColorMap:
+            colormap = {}
+            count = 2 ** self.ifds[0].BitsPerSample.value
+
+            nodata_val = None
+            if self.has_alpha or self.nodata is not None:
+                nodata_val = 0 if self.has_alpha else self.nodata
+
+            transform = lambda val: int((val / 65535) * 255)
+            for idx in range(count):
+                color = [transform(self.ifds[0].ColorMap.value[idx + i * count]) for i in range(3)]
+                if nodata_val is not None:
+                    color.append(0 if idx == nodata_val else 255)
+                colormap[idx] = tuple(color)
+            return colormap
+        return None
+
+    @property
+    def color_interp(self):
+        """
+        https://gdal.org/user/raster_data_model.html#raster-band
+        https://trac.osgeo.org/gdal/ticket/4547#comment:1
+        """
+        photometric = self.photometric
+        if photometric == "rgb":
+            interp = [ColorInterp.red, ColorInterp.green, ColorInterp.blue]
+            if self.has_alpha:
+                interp.append(ColorInterp.alpha)
+        elif photometric == "minisblack" or photometric == "miniswhite":
+            interp = [ColorInterp.gray]
+        elif photometric == "palette":
+            interp = [ColorInterp.palette]
+        elif photometric == "cmyk":
+            interp = [ColorInterp.cyan, ColorInterp.magenta, ColorInterp.yellow, ColorInterp.black]
+        elif photometric == "ycbcr":
+            interp = [ColorInterp.red, ColorInterp.green, ColorInterp.blue]
+        elif photometric == "cielab" or photometric == "icclab" or photometric == "itulab":
+            interp = [ColorInterp.lightness, ColorInterp.lightness, ColorInterp.lightness]
+        else:
+            interp = [ColorInterp.undefined for _ in range(self.profile['count'])]
+        return interp
+
+
+
     @property
     def has_alpha(self) -> bool:
         """Check if the image has an alpha band"""

aiocogeo/constants.py

@@ -14,14 +14,34 @@
 
 INTERLEAVE = {1: "pixel", 2: "band"}
 
-# https://www.awaresystems.be/imaging/tiff/tifftags/photometricinterpretation.html
+class ColorInterp(enum.IntEnum):
+    """https://github.com/mapbox/rasterio/blob/master/rasterio/enums.py#L6-L25"""
+    undefined = 0
+    gray = 1
+    grey = 1
+    palette = 2
+    red = 3
+    green = 4
+    blue = 5
+    alpha = 6
+    hue = 7
+    saturation = 8
+    lightness = 9
+    cyan = 10
+    magenta = 11
+    yellow = 12
+    black = 13
+    Y = 14
+    Cb = 15
+    Cr = 16
+
 PHOTOMETRIC = {
     0: "miniswhite",
     1: "minisblack",
     2: "rgb",
     3: "palette",
     4: "mask",
-    5: "separated",
+    5: "cmyk",
     6: "ycbcr",
     8: "cielab",
     9: "icclab",
@@ -128,6 +148,7 @@
     277: "SamplesPerPixel",
     284: "PlanarConfiguration",
     317: "Predictor",
+    320: "ColorMap",
     322: "TileWidth",
     323: "TileHeight",
     324: "TileOffsets",
@@ -146,4 +167,4 @@ class MaskFlags(enum.IntEnum):
     all_valid = 1
     per_dataset = 2
     alpha = 4
-    nodata = 8
+    nodata = 8

aiocogeo/ifd.py

@@ -72,6 +72,7 @@ class OptionalTags:
     Predictor: Tag = None
     JPEGTables: Tag = None
     ExtraSamples: Tag = None
+    ColorMap: Tag = None
 
     # GeoTiff
     GeoKeyDirectoryTag: Tag = None

aiocogeo/utils.py

@@ -1,6 +1,6 @@
 import asyncio
 from functools import partial
-from typing import Any, Callable
+from typing import Any, Callable, List
 
 
 async def run_in_background(
@@ -17,4 +17,10 @@ async def run_in_background(
     """
     loop = asyncio.get_event_loop()
     func = partial(func, *args, **kwargs)
-    return await loop.run_in_executor(None, func)
+    return await loop.run_in_executor(None, func)
+
+
+def chunks(lst: List, n: int):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i:i + n]

tests/test_cog_reader.py

@@ -38,6 +38,7 @@ async def test_cog_metadata(infile, create_cog_reader):
             cog_profile.pop("photometric", None)
             rio_profile.pop("photometric", None)
 
+            assert [member.value for member in ds.colorinterp] == [member.value for member in cog.color_interp]
             assert rio_profile == cog_profile
             assert ds.overviews(1) == cog.overviews
 
@@ -478,6 +479,17 @@ async def test_read_not_in_bounds(create_cog_reader, infile):
             await cog.read(bounds=bounds, shape=(256, 256))
 
 
+@pytest.mark.asyncio
+async def test_cog_palette(create_cog_reader):
+    infile = "https://async-cog-reader-test-data.s3.amazonaws.com/cog_cmap.tif"
+    async with create_cog_reader(infile) as cog:
+        with rasterio.open(infile) as ds:
+            cog_interp = cog.color_interp
+            rio_interp = ds.colorinterp
+            assert cog_interp[0].value == rio_interp[0].value
+            assert cog.colormap == ds.colormap(1)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "width,height", [(500, 500), (1000, 1000), (5000, 5000), (10000, 10000)]