From c435d06e55fbed71a5cfbcb1ba5375986930be1a Mon Sep 17 00:00:00 2001
From: tebben <ebben.tim@gmail.com>
Date: Mon, 22 Jan 2024 13:03:44 +0100
Subject: [PATCH] Block loading unsafe request when not enough overviews and
 show warning

---
 ctod/core/cog/cog_reader.py                   | 41 +++++++++++++++++--
 ctod/core/layer.py                            |  6 ++-
 .../terrain_generator_quantized_mesh_grid.py  |  2 +-
 3 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/ctod/core/cog/cog_reader.py b/ctod/core/cog/cog_reader.py
index cbb708b..383ea87 100644
--- a/ctod/core/cog/cog_reader.py
+++ b/ctod/core/cog/cog_reader.py
@@ -1,6 +1,7 @@
 import time
-
-from morecantile import TileMatrixSet
+import math
+import logging
+from morecantile import TileMatrixSet, Tile
 from rio_tiler.io import Reader
 from rio_tiler.models import ImageData
 
@@ -35,9 +36,14 @@ def download_tile(self, x: int, y: int, z: int, resampling_method="bilinear") ->
             resampling_method (str, optional): RasterIO resampling algorithm. Defaults to "bilinear".
 
         Returns:
-            ImageData: _description_
+            ImageData: Image data from the Cloud Optimized GeoTIFF.
         """
         
+        is_safe = self._check_is_safe(z, x, y)
+        if not is_safe:
+            logging.warning(f"Skipping {self.cog} {z,x,y}, tile is not safe to load, generate more overviews")
+            return None
+        
         try:
 
             image_data = self.rio_reader.tile(tile_z=z, tile_x=x, tile_y=y, resampling_method=resampling_method)
@@ -52,6 +58,30 @@ def download_tile(self, x: int, y: int, z: int, resampling_method="bilinear") ->
         except Exception:
             return None
     
+    def _check_is_safe(self, z: int, x: int, y: int) -> bool:
+        """Check if it is safe to load the tile. 
+        When there are not enough overviews there will be a lot of request to load
+        a tile at a low zoom level. This will cause long load times and high memory usage.
+
+        ToDo: This is an estimation, it is not 100% accurate.        
+        
+        Args:
+            z (int): 
+            x (int): 
+            y (int): 
+
+        Returns:
+            bool: Is it safe to load the tile
+        """
+        
+        tile_bounds = self.tms.xy_bounds(Tile(x=x, y=y, z=z))        
+        tile_wgs = tile_bounds.right - tile_bounds.left
+        #tile_wgs_with_clip = min(tile_bounds.right, self.dataset_bounds.right) - min(tile_bounds.left, dataset_bounds.left)        
+        tile_pixels_needed = tile_wgs * self.pixels_per_wgs        
+        needed_tiles = math.ceil(tile_pixels_needed / self.pixels_per_tile_downsampled)
+        
+        return needed_tiles <= 1        
+            
     def return_reader(self):
         """Done with the reader, return it to the pool."""
         
@@ -62,6 +92,11 @@ def _set_rio_reader(self):
         """Get the reader for the COG."""
         
         self.rio_reader = Reader(self.cog, tms=self.tms)
+        self.dataset_bounds = self.rio_reader.info().bounds
+        self.dataset_width = self.rio_reader.dataset.width
+        self.dataset_wgs_width = self.dataset_bounds.right - self.dataset_bounds.left
+        self.pixels_per_wgs = self.dataset_width / self.dataset_wgs_width
+        self.pixels_per_tile_downsampled = 256 * max(self.rio_reader.dataset.overviews(1))
         
     def _set_nodata_value(self):
         """Set the nodata value for the reader."""
diff --git a/ctod/core/layer.py b/ctod/core/layer.py
index 9bb40e3..80e8a08 100644
--- a/ctod/core/layer.py
+++ b/ctod/core/layer.py
@@ -8,9 +8,13 @@ def generate_layer_json(tms, geotiff_path, max_zoom=20):
     """
     
     with COGReader(geotiff_path) as src:
+        
         bounds = src.geographic_bounds
         
-        # Zoom 0 always needs to be 0,0,1,0
+        # Cesium always expects all tiles at zoom 0 (startX: 0, endX: 1)
+        # With the function available_tiles it is likely it only will return 
+        # one tile: startX: 1, endX: 1 for example. So here we skip generating
+        # the first zoom level
         available_tiles = [
             [{"startX": 0, "startY": 0, "endX": 1, "endY": 0}]
         ]
diff --git a/ctod/core/terrain/generator/terrain_generator_quantized_mesh_grid.py b/ctod/core/terrain/generator/terrain_generator_quantized_mesh_grid.py
index e9facb6..a6dcae4 100644
--- a/ctod/core/terrain/generator/terrain_generator_quantized_mesh_grid.py
+++ b/ctod/core/terrain/generator/terrain_generator_quantized_mesh_grid.py
@@ -176,7 +176,7 @@ def _get_vertice_condition(self, vertices, direction) -> np.ndarray:
         if direction == Direction.NORTH:
             vertices[:, 1] -= tile_size
             return vertices[:, 1] == 0             
-        if direction == Direction.NORTHEAST:
+        elif direction == Direction.NORTHEAST:
             vertices[:, 0] -= tile_size
             vertices[:, 1] -= tile_size
             return (vertices[:, 0] == 0) & (vertices[:, 1] == 0)