tilegrab 1.1.0__py3-none-any.whl → 1.2.0b2__py3-none-any.whl

tilegrab/images.py

@@ -1,35 +1,44 @@
+from collections import defaultdict
+from functools import cache
 import logging
 from dataclasses import dataclass
-from io import BytesIO
 from pathlib import Path, PosixPath, WindowsPath
-import re
-from typing import Any, List, Optional, Tuple, Union
-from PIL import Image as PLIImage
+from typing import Any, List, Union
+from PIL import Image as PILImage
 from box import Box
 from tilegrab.tiles import Tile, TileCollection
 import os
+import numpy as np
+from numpy.lib.stride_tricks import sliding_window_view
 
 logger = logging.getLogger(__name__)
 
+
 class ExportType:
     PNG: int = 1
     JPG: int = 2
     TIFF: int = 3
 
+
 @dataclass
 class TileImage:
     width: int = 256
     height: int = 256
 
     def __init__(self, tile: Tile, image: Union[bytes, bytearray]) -> None:
+        from io import BytesIO
+
         self._tile = tile
         try:
-            self._img = PLIImage.open(BytesIO(image))
+            self._img = PILImage.open(BytesIO(image))
             logger.debug(f"TileImage created for z={tile.z},x={tile.x},y={tile.y}")
         except Exception as e:
-            logger.error(f"Failed to open image for tile z={tile.z},x={tile.x},y={tile.y}", exc_info=True)
+            logger.error(
+                f"Failed to open image for tile z={tile.z},x={tile.x},y={tile.y}",
+                exc_info=True,
+            )
             raise
-        
+
         self._path: Union[Path, None] = None
         self._ext: str = self._get_image_type(image)
 
@@ -40,23 +49,23 @@ class TileImage:
         b = bytes(data)
 
         # PNG: 8 bytes
-        if b.startswith(b'\x89PNG\r\n\x1a\n'):
+        if b.startswith(b"\x89PNG\r\n\x1a\n"):
             logger.debug(f"Image detected as PNG")
-            return 'png'
+            return "png"
 
         # JPEG / JPG: files start with FF D8 and end with FF D9
-        if len(b) >= 2 and b[0:2] == b'\xff\xd8':
+        if len(b) >= 2 and b[0:2] == b"\xff\xd8":
             logger.debug(f"Image detected as JPG")
-            return 'jpg'
+            return "jpg"
 
         # BMP: starts with 'BM' (0x42 0x4D)
-        if len(b) >= 2 and b[0:2] == b'BM':
+        if len(b) >= 2 and b[0:2] == b"BM":
             logger.debug(f"Image detected as BMP")
-            return 'bmp'
+            return "bmp"
 
         logger.warning(f"Unknown image format, defaulting to PNG")
         return "png"
-    
+
     def save(self):
         try:
             self._img.save(self.path)
@@ -68,16 +77,16 @@ class TileImage:
     @property
     def name(self) -> str:
         return f"{self._tile.z}_{self._tile.x}_{self._tile.y}.{self._ext}"
-    
+
     @property
     def tile(self) -> Tile:
         return self._tile
-    
+
     @property
-    def image(self) -> PLIImage.Image:
+    def image(self) -> PILImage.Image:
         self._img.load()
         return self._img
-    
+
     @property
     def path(self) -> Path:
         if self._path is None:
@@ -133,15 +142,68 @@ class TileImage:
         return self._tile.url
 
 
+WEB_MERCATOR_EXTENT = 20037508.342789244
+EPSG = 3857
+
+
 class TileImageCollection:
     images: List[TileImage] = []
     width: int = 0
     height: int = 0
 
-    def load(self, tile_collection:TileCollection):
+    def __init__(self, path: Union[Path, str]) -> None:
+        self.path = Path(path)
+        logger.info(f"TileImageCollection initialized at {self.path}")
+
+    def __len__(self):
+        return sum(1 for _ in self)
+
+    def __iter__(self):
+        for i in self.images:
+            yield i
+
+    def __repr__(self) -> str:
+        return f"ImageCollection; len={len(self)}"
+
+    def _update_collection_dim(self):
+        if not self.images:
+            logger.warning("Attempting to update collection dimensions with no images")
+            return
+
+        x = [img.tile.x for img in self.images]
+        y = [img.tile.y for img in self.images]
+        minx, maxx = min(x), max(x)
+        miny, maxy = min(y), max(y)
+        self.minx, self.maxx = minx, maxx
+        self.miny, self.maxy = miny, maxy
+        logger.debug(f"Tile range x=({self.minx}, {self.maxx}); y=({self.miny}, {self.maxy})")
+
+        self.width = int((maxx - minx + 1) * self.images[0].width)
+        self.height = int((maxy - miny + 1) * self.images[0].height)
+       
+        
+        logger.info(f"Collection dimensions calculated: {self.width}x{self.height}")
+
+    def mosaic_bounds(self, x_min, y_min, x_max, y_max, z):
+        n = 2**z
+        tile_size_m = 2 * WEB_MERCATOR_EXTENT / n
+
+        xmin = (WEB_MERCATOR_EXTENT * -1) + x_min * tile_size_m
+        xmax = (WEB_MERCATOR_EXTENT * -1) + (x_max + 1) * tile_size_m
+
+        ymax = WEB_MERCATOR_EXTENT - y_min * tile_size_m
+        ymin = WEB_MERCATOR_EXTENT - (y_max + 1) * tile_size_m
+
+        return xmin, ymin, xmax, ymax
+
+    def load(self, tile_collection: TileCollection):
+        import re
+
         logger.info("Start loading saved ImageTiles")
-        pat = re.compile(r'^([0-9]+)_([0-9]+)_([0-9]+)\.[A-Za-z0-9]+$')
+        pat = re.compile(r"^([0-9]+)_([0-9]+)_([0-9]+)\.[A-Za-z0-9]+$")
         image_col = [p for p in self.path.glob(f"*.*") if p.is_file()]
+        self.zoom = tile_collection.to_list[0].z
+        logger.info(f"Found {len(image_col)} images at {self.path}")
 
         for tile in tile_collection.to_list:
             found_matching_image = False
@@ -154,66 +216,125 @@ class TileImageCollection:
 
                     if tile.x == x and tile.y == y and tile.z == z:
                         logger.debug(f"Processing ImageTile x={x} y={y} z={z}")
-                        with open(image_path, 'rb') as f:
+                        with open(image_path, "rb") as f:
                             tile_image = TileImage(tile, f.read())
                             tile_image.path = image_path
                             self.images.append(tile_image)
                             found_matching_image = True
                             continue
 
-            if not found_matching_image: logger.warning(f"Missing ImageTile x={tile.x} y={tile.y} z={tile.z}")
+            if not found_matching_image:
+                logger.warning(f"Missing ImageTile x={tile.x} y={tile.y} z={tile.z}")
         
+        logger.info(f"{len(self.images)} images loaded, {len(image_col) - len(self.images)} skipped")
+
+        self._update_collection_dim()
+
     def append(self, img: TileImage):
         img.path = os.path.join(self.path, img.name)
         self.images.append(img)
         logger.debug(f"Image appended to collection: {img.name}")
         img.save()
 
-    def __init__(self, path: Union[Path, str]) -> None:
-        self.path = Path(path)
-        logger.info(f"TileImageCollection initialized at {self.path}")
+        self.zoom = img.tile.z
 
-    def __len__(self):
-        return sum(1 for _ in self)
-
-    def __iter__(self):
-        for i in self.images:
-            yield i
-
-    def __repr__(self) -> str:
-        return f"ImageCollection; len={len(self)}"
-
-    def _update_collection_dim(self):
-        if not self.images:
-            logger.warning("Attempting to update collection dimensions with no images")
-            return
-        
-        x = [img.tile.x for img in self.images]
-        y = [img.tile.y for img in self.images]
-        minx, maxx = min(x), max(x)
-        miny, maxy = min(y), max(y)
-
-        self.width = int((maxx - minx + 1) * self.images[0].width)
-        self.height = int((maxy - miny + 1) * self.images[0].height)
-        logger.info(f"Collection dimensions calculated: {self.width}x{self.height}")
-
-    def mosaic(self):
-        logger.info("Starting mosaic creation")
+    @cache
+    def _create_mosaic(self) -> PILImage.Image:
+        logger.info("Start mosaicing TileImageCollection")
         self._update_collection_dim()
 
-        logger.info(f"Mosaicking {len(self.images)} images into {self.width}x{self.height}")
-        merged_image = PLIImage.new("RGB", (self.width, self.height))
+        logger.info(
+            f"Mosaicking {len(self.images)} images into {self.width}x{self.height}"
+        )
+        merged_image = PILImage.new("RGB", (self.width, self.height))
 
         for image in self.images:
             px = int((image.position.x) * image.width)
             py = int((image.position.y) * image.height)
             logger.debug(f"Pasting image at position ({px}, {py}): {image.name}")
             merged_image.paste(image.image, (px, py))
+        
+        return merged_image
 
-        output_path = "merged_output.png"
-        merged_image.save(output_path)
-        logger.info(f"Mosaic saved to {output_path}")
+    def mosaic(self, export_types: List[int]):
+        
+        merged_image = self._create_mosaic()
+        # TODO: fix this monkey patch
+        if ExportType.TIFF in export_types:
+            output_path = "mosaic.tiff"
+            import numpy as np
+            from rasterio.transform import from_bounds
+            import rasterio
+
+            data = np.array(merged_image)
+            data = data.transpose(2, 0, 1)
+
+            width_px, height_px = merged_image.size
+            xmin, ymin, xmax, ymax = self.mosaic_bounds(
+                self.minx, self.miny, self.maxx, self.maxy, self.images[0].tile.z
+            )
+
+            transform = from_bounds(xmin, ymin, xmax, ymax, width_px, height_px)
+
+            with rasterio.open(
+                output_path,
+                "w",
+                driver="GTiff",
+                height=height_px,
+                width=width_px,
+                count=data.shape[0],
+                dtype=data.dtype,
+                crs=f"EPSG:{EPSG}",
+                transform=transform,
+            ) as dst:
+                dst.write(data)
+
+            logger.info(f"Mosaic saved to {output_path}")
+        if ExportType.PNG in export_types:
+            output_path = "mosaic.png"
+            merged_image.save(output_path)
+            logger.info(f"Mosaic saved to {output_path}")
+        if ExportType.JPG in export_types:
+            output_path = "mosaic.jpg"
+            merged_image.save(output_path)
+            logger.info(f"Mosaic saved to {output_path}")
+
+    def group(self, width: int, height: int, export_types:List[int], overlap: bool):
+        if width <= 0 or height <= 0:
+            raise ValueError("width and height must be positive integers")
+        
+        out_dir = "grouped_tiles"
+        os.makedirs(out_dir, exist_ok=True)
+        logger.info(f"Start grouping {len(self.images)} TileImages into {width}x{height} groups")
+
+        merged_image= self._create_mosaic()
+        merged_image_arr = np.asanyarray(merged_image)
+
+        image_px_width = 256
+        image_px_height = 256
+        kh, kw = image_px_height*height, image_px_width*width
+        view = sliding_window_view(merged_image_arr, (kh, kw, merged_image_arr.shape[2]))
+        
+        view = view[..., 0, :, :, :]
+        OH, OW = view.shape[:2]
+        logger.info(f"Saving groups into ./{out_dir}")
+        for i in range(0, OH, height*image_px_height):
+            for j in range(0, OW, width*image_px_width):
+                patch = view[i, j]
+                if not (patch == 0).all():
+                    img = PILImage.fromarray(patch)
+                    if ExportType.PNG in export_types:
+                        logger.debug(f"Saving group {i}{j} as {i}{j}.png")
+                        img.save(os.path.join(out_dir, f"{i}{j}.png"))
+                    if ExportType.JPG in export_types:
+                        logger.debug(f"Saving group {i}{j} as {i}{j}.jpg")
+                        img.save(os.path.join(out_dir, f"{i}{j}.jpg"))
+                    if ExportType.TIFF in export_types:
+                        logger.error(f"TIFF exports for tile groups not implemented yet")
+                        raise NotImplementedError
+                else:
+                    logger.debug(f"Skip no-data group {i}{j}")
 
     def export_collection(self, type: ExportType):
         logger.info(f"Exporting collection as type {type}")
-        pass
+        raise NotImplementedError

tilegrab/logs.py

@@ -0,0 +1,84 @@
+import logging
+import sys
+
+# Normal colors
+BLACK   = "\033[30m"
+RED     = "\033[31m"
+GREEN   = "\033[32m"
+YELLOW  = "\033[33m"
+BLUE    = "\033[34m"
+MAGENTA = "\033[35m"
+CYAN    = "\033[36m"
+GRAY    = "\033[90m"
+WHITE  = "\033[37m"
+
+# Bright colors
+BBLACK   = "\033[90m"
+BRED     = "\033[91m"
+BGREEN   = "\033[92m"
+BYELLOW  = "\033[93m"
+BBLUE    = "\033[94m"
+BMAGENTA = "\033[95m"
+BCYAN    = "\033[96m"
+BGRAY    = "\033[97m"
+BWHITE = "\033[97m"
+
+RESET = "\033[0m"
+
+
+class CLILogFormatter(logging.Formatter):
+    NAME_WIDTH = 20
+    LEVEL_MAP = {
+        logging.CRITICAL: f'{RED}‼ {RESET}',
+        logging.ERROR:    f'{RED}✖ {RESET}',
+        logging.WARNING:  f'{YELLOW}⚠ {RESET}',
+        logging.INFO:     f'{BLUE}• {RESET}',
+        logging.DEBUG:    f'{GRAY}· {RESET}',
+        logging.NOTSET:   f'{CYAN}- {RESET}',
+    }
+
+    def __init__(self, fmt=None):
+        super().__init__(fmt or '   %(level_icon)s %(message)s')
+
+    def format(self, record):
+        record.level_icon = self.LEVEL_MAP.get(record.levelno, '?')
+        short = record.name.rsplit('.', 1)[-1]
+        record.short_name = f"{short:<{self.NAME_WIDTH}}"
+
+        return super().format(record)
+
+class FileLogFormatter(logging.Formatter):
+    def __init__(self):
+        super().__init__(
+            '%(asctime)s %(levelname)s %(name)s - %(message)s'
+        )
+
+def create_cli_handler(level=logging.INFO):
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(level)
+    handler.setFormatter(CLILogFormatter())
+    return handler
+
+def create_file_handler(path="tilegrab.log", level=logging.DEBUG):
+    handler = logging.FileHandler(path)
+    handler.setLevel(level)
+    handler.setFormatter(FileLogFormatter())
+    return handler
+
+def setup_logging(
+    enable_cli=True,
+    enable_file=True,
+    level=logging.INFO
+):
+    handlers = []
+
+    if enable_cli:
+        handlers.append(create_cli_handler(level))
+
+    if enable_file:
+        handlers.append(create_file_handler(level=level))
+
+    logging.basicConfig(
+        level=level,
+        handlers=handlers,
+    )

tilegrab/sources.py

@@ -4,8 +4,10 @@ from typing import Dict, Optional
 logger = logging.getLogger(__name__)
 
 class TileSource:
-    URL_TEMPLATE = ""
-    name = None
+    url_template: str
+    name: str
+    description: str
+    output_dir: str
 
     def __init__(
         self, 
@@ -17,7 +19,7 @@ class TileSource:
         logger.debug(f"Initializing TileSource: {self.name}, has_api_key={api_key is not None}")
 
     def get_url(self, z: int, x: int, y: int) -> str:
-        url = self.URL_TEMPLATE.format(x=x, y=y, z=z)
+        url = self.url_template.format(x=x, y=y, z=z)
         logger.debug(f"Generated URL for {self.name}: z={z}, x={x}, y={y}")
         return url
 
@@ -33,30 +35,34 @@ class TileSource:
 
 
 class GoogleSat(TileSource):
+    output_dir = "ggl_sat"
     name = "GoogleSat"
     description = "Google satellite imageries"
-    URL_TEMPLATE = "https://mt1.google.com/vt/lyrs=s&x={x}&y={y}&z={z}"
+    url_template = "https://mt1.google.com/vt/lyrs=s&x={x}&y={y}&z={z}"
     
 
 class OSM(TileSource):
+    output_dir = "osm"
     name = "OSM"
     description = "OpenStreetMap imageries"
-    URL_TEMPLATE = "https://tile.openstreetmap.org/{z}/{x}/{y}.png"
+    url_template = "https://tile.openstreetmap.org/{z}/{x}/{y}.png"
 
 class ESRIWorldImagery(TileSource):
+    output_dir = "esri_world"
     name = "ESRIWorldImagery"
     description = "ESRI satellite imageries"
-    URL_TEMPLATE = "https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}"
+    url_template = "https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}"
 
 class Nearmap(TileSource):
+    output_dir = "nearmap_sat"
     name = "NearmapSat"
     description = "Nearmap satellite imageries"
-    URL_TEMPLATE = "https://api.nearmap.com/tiles/v3/Vert/{z}/{x}/{y}.png?apikey={token}"
+    url_template = "https://api.nearmap.com/tiles/v3/Vert/{z}/{x}/{y}.png?apikey={token}"
 
     def get_url(self, z: int, x: int, y: int) -> str:
         if not self.api_key:
             logger.error("Nearmap API key is required but not provided")
             raise AssertionError("API key required for Nearmap")
-        url = self.URL_TEMPLATE.format(x=x, y=y, z=z, token=self.api_key)
+        url = self.url_template.format(x=x, y=y, z=z, token=self.api_key)
         logger.debug(f"Generated Nearmap URL: z={z}, x={x}, y={y}")
         return url