Use disk-backed memmap for tile stitching to avoid OOM

Author: fcladera

scripts/tile_composer.py

@@ -1,5 +1,7 @@
 import math
+import os
 import requests
+import numpy as np
 from PIL import Image
 from io import BytesIO
 from pathlib import Path
@@ -96,10 +98,15 @@ def download_and_save(lat_max, lat_min, lon_min, lon_max, zoom, image_name):
     num_y = y_max - y_min + 1
 
 
-    # Replace it with these:
-    stitched_image = None
     tile_size = None
     headers = {'User-Agent': 'Mozilla/5.0'}
+    mmap_path = f"{image_name}_canvas.dat"
+
+    # Clean up stale memmap from a previous run
+    if os.path.exists(mmap_path):
+        os.remove(mmap_path)
+
+    canvas = None
 
     image_min_lat = 90.0
     image_min_long = 180.0
@@ -108,23 +115,24 @@ def download_and_save(lat_max, lat_min, lon_min, lon_max, zoom, image_name):
 
     for i, x in enumerate(range(x_min, x_max + 1)):
         for j, y in enumerate(range(y_min, y_max + 1)):
-            url = get_satellite_url(x, y, zoom) # Add provider="mapbox" etc. if using your updated function
+            url = get_satellite_url(x, y, zoom)
             response = requests.get(url, headers=headers)
-            
+
             if response.status_code == 200:
                 tile = Image.open(BytesIO(response.content))
-                
-                # --- LAZY CANVAS INITIALIZATION ---
-                # If this is the first valid tile, measure it and build the canvas
-                if stitched_image is None:
-                    tile_size = tile.width  # This will automatically be 256 or 512
+
+                # Lazy init: measure first tile and create disk-backed canvas
+                if canvas is None:
+                    tile_size = tile.width
                     print(f"Detected tile size: {tile_size}x{tile_size}px")
-                    stitched_image = Image.new('RGB', (num_x * tile_size, num_y * tile_size))
-                
-                # Paste using the dynamic tile_size instead of 256 or 512
-                stitched_image.paste(tile, (i * tile_size, j * tile_size))
+                    print(f"Canvas: {num_x * tile_size}x{num_y * tile_size}px (memmap)")
+                    canvas = np.memmap(mmap_path, dtype=np.uint8, mode='w+',
+                                       shape=(num_y * tile_size, num_x * tile_size, 3))
+
+                tile_arr = np.array(tile.convert('RGB'))
+                canvas[j * tile_size:(j + 1) * tile_size,
+                       i * tile_size:(i + 1) * tile_size] = tile_arr
 
-                # Get corners
                 corners = get_tile_corners(x, y, zoom)
                 for _, val in corners.items():
                     lat = val[0]
@@ -140,34 +148,32 @@ def download_and_save(lat_max, lat_min, lon_min, lon_max, zoom, image_name):
             else:
                 print(f"Error: Could not download tile {x},{y}")
 
-    # Failsafe in case the entire grid failed to download
-    if stitched_image is None:
+    if canvas is None:
         print("Error: Could not download any tiles. Canvas was never created.")
         return None
-   
-    # Determine the global pixel coordinate of the top-left corner of your stitched image
-    # (Which is the NW corner of tile x_min, y_min)
+
     ref_pixel_x = x_min * tile_size
     ref_pixel_y = y_min * tile_size
 
-    # Get the global pixel coordinates for your requested bounding box
     upper_left_x, upper_left_y = latlon_to_pixel(lat_max, lon_min, zoom, tile_size)
     lower_right_x, lower_right_y = latlon_to_pixel(lat_min, lon_max, zoom, tile_size)
 
-    # Calculate the crop boundaries relative to the stitched image (0,0)
-    left = upper_left_x - ref_pixel_x
-    top = upper_left_y - ref_pixel_y
-    right = lower_right_x - ref_pixel_x
-    bottom = lower_right_y - ref_pixel_y
+    left = int(upper_left_x - ref_pixel_x)
+    top = int(upper_left_y - ref_pixel_y)
+    right = int(lower_right_x - ref_pixel_x)
+    bottom = int(lower_right_y - ref_pixel_y)
 
-    # Crop the output image
     print(f"Cropping image to: {left, top, right, bottom}")
-    stitched_image = stitched_image.crop((left, top, right, bottom))
+    cropped = Image.fromarray(canvas[top:bottom, left:right])
 
     output_path = Path(f"{image_name}.jpg")
     output_path.parent.mkdir(parents=True, exist_ok=True)
-    stitched_image.save(output_path, "JPEG", quality=90)
-    
+    cropped.save(output_path, "JPEG", quality=90)
+
+    # Clean up memmap file
+    del canvas
+    os.remove(mmap_path)
+
     print(f"Rectangular satellite image saved as '{image_name}.jpg'")
     print(f"Bounds: {image_min_lat},{image_min_long} {image_max_lat},{image_max_long}")
     return str(output_path)