Merge pull request #35 from ucam-eo/fetch_embeddings_lazy

Fetch embeddings lazily and embeddings count

Author: avsm

geotessera/core.py

@@ -6,7 +6,7 @@
 """
 
 from pathlib import Path
-from typing import Union, List, Tuple, Optional, Dict
+from typing import Union, List, Tuple, Optional, Dict, Generator
 import json
 import numpy as np
 
@@ -59,21 +59,36 @@ def version(self) -> str:
         """Get the GeoTessera library version."""
         return __version__
 
+    def embeddings_count(self, bbox: Tuple[float, float, float, float], year: int = 2024) -> int:
+            """Get total number of embedding tiles within a bounding box.
+
+            Args:
+                bbox: Bounding box as (min_lon, min_lat, max_lon, max_lat)
+                year: Year of embeddings to consider
+
+            Returns:
+                Total number of tiles in the bounding box
+            """
+            tiles = self.registry.load_blocks_for_region(bbox, year)
+            return len(tiles)
+
+    # returns a generator
     def fetch_embeddings(
         self,
         bbox: Tuple[float, float, float, float],
         year: int = 2024,
         progress_callback: Optional[callable] = None,
-    ) -> List[Tuple[float, float, np.ndarray, object, object]]:
-        """Fetch all embedding tiles within a bounding box with CRS information.
+    ) -> Generator[Tuple[float, float, np.ndarray, object, object], None, None]:
+        """Lazily fetches all embedding tiles within a bounding box with CRS information.
+        Use as a generator to process tiles one at a time in a memory-efficient manner.
 
         Args:
             bbox: Bounding box as (min_lon, min_lat, max_lon, max_lat)
             year: Year of embeddings to download
             progress_callback: Optional callback function(current, total) for progress tracking
 
         Returns:
-            List of (tile_lon, tile_lat, embedding_array, crs, transform) tuples where:
+            Generator of (tile_lon, tile_lat, embedding_array, crs, transform) tuples where:
             - tile_lon: Tile center longitude
             - tile_lat: Tile center latitude
             - embedding_array: shape (H, W, 128) with dequantized values
@@ -84,7 +99,6 @@ def fetch_embeddings(
         tiles_to_download = self.registry.load_blocks_for_region(bbox, year)
 
         # Download each tile with progress tracking
-        results = []
         total_tiles = len(tiles_to_download)
 
         for i, (tile_lon, tile_lat) in enumerate(tiles_to_download):
@@ -108,7 +122,8 @@ def tile_progress_callback(
                 embedding, crs, transform = self.fetch_embedding(
                     tile_lon, tile_lat, year, tile_progress_callback
                 )
-                results.append((tile_lon, tile_lat, embedding, crs, transform))
+
+                yield tile_lon, tile_lat, embedding, crs, transform
 
                 # Update progress for completed tile
                 if progress_callback:
@@ -130,7 +145,7 @@ def tile_progress_callback(
                     )
                 continue
 
-        return results
+        return None 
 
     def fetch_embedding(
         self,
@@ -601,7 +616,7 @@ def merge_geotiffs_to_mosaic(
                 if progress_callback:
                     progress_callback(i, total_files * 2 + 2, f"Reprojecting file {i+1}/{total_files}...")
 
-                result_file, error = self._reproject_geotiff_file(args)
+                _, error = self._reproject_geotiff_file(args)
 
                 if error:
                     failed_files.append((geotiff_paths[i], error))