perf(heightmaps): Use in-memory warping for heightmap export
Improved performance and reduced disk I/O by using GDAL's MEM driver for intermediate warp operations. The warped dataset is now passed directly to the translate step without being written to a temporary file.
This commit is contained in:
64
tests/test_heightmaps_optimization.py
Normal file
64
tests/test_heightmaps_optimization.py
Normal file
@@ -0,0 +1,64 @@
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import os
|
||||
from geodata_pipeline.heightmaps import export_heightmaps
|
||||
from geodata_pipeline.config import Config
|
||||
|
||||
class TestHeightmapsOptimization(unittest.TestCase):
|
||||
@patch("geodata_pipeline.heightmaps.gdal.Warp")
|
||||
@patch("geodata_pipeline.heightmaps.gdal.Translate")
|
||||
@patch("geodata_pipeline.heightmaps.open_dataset")
|
||||
@patch("geodata_pipeline.heightmaps.build_vrt")
|
||||
@patch("geodata_pipeline.heightmaps.ensure_dir")
|
||||
@patch("geodata_pipeline.heightmaps.ensure_parent")
|
||||
@patch("geodata_pipeline.heightmaps._compute_ds_minmax")
|
||||
@patch("geodata_pipeline.heightmaps.cleanup_aux_files")
|
||||
@patch("geodata_pipeline.heightmaps.glob.glob")
|
||||
@patch("builtins.open")
|
||||
def test_export_heightmaps_in_memory_warp(self, mock_open_file, mock_glob, mock_cleanup, mock_minmax, mock_ensure_p, mock_ensure_d, mock_build_vrt, mock_open_ds, mock_translate, mock_warp):
|
||||
# Setup mocks
|
||||
mock_glob.return_value = ["tile1.tif"]
|
||||
mock_open_ds.return_value.GetRasterBand.return_value.ComputeRasterMinMax.return_value = (0, 100)
|
||||
mock_minmax.return_value = (10, 90, 100) # min, max, count
|
||||
|
||||
# Mock VRT dataset
|
||||
mock_ds = MagicMock()
|
||||
mock_open_ds.return_value = mock_ds
|
||||
mock_ds.GetRasterBand.return_value.ComputeRasterMinMax.return_value = (0, 100)
|
||||
|
||||
# Mock individual tile dataset for GT
|
||||
mock_tds = MagicMock()
|
||||
mock_tds.GetGeoTransform.return_value = (1000, 1, 0, 2000, 0, -1)
|
||||
mock_tds.RasterXSize = 1000
|
||||
mock_tds.RasterYSize = 1000
|
||||
|
||||
# Side effect for open_dataset: first VRT, then tile1.tif
|
||||
mock_open_ds.side_effect = [mock_ds, mock_tds]
|
||||
|
||||
cfg = Config.default()
|
||||
cfg.raw.dgm1_dir = "raw/dgm1"
|
||||
cfg.work.work_dir = "work"
|
||||
cfg.export.heightmap_dir = "export"
|
||||
cfg.export.manifest_path = "manifest.csv"
|
||||
|
||||
# Mock Warp to return a memory dataset
|
||||
mock_mem_ds = MagicMock()
|
||||
mock_warp.return_value = mock_mem_ds
|
||||
|
||||
export_heightmaps(cfg)
|
||||
|
||||
# Verify Warp was called with empty string or MEM driver destination
|
||||
warp_args, warp_kwargs = mock_warp.call_args
|
||||
# Warp(destNameOrDestDS, srcDSOrSrcDSTab, ...)
|
||||
dest = warp_args[0]
|
||||
# In the original code, dest was a path. We want to check it's changed to "" for MEM.
|
||||
# But wait, my refactor will use "" for destName.
|
||||
|
||||
# Verify Translate used the warped dataset
|
||||
translate_args, translate_kwargs = mock_translate.call_args
|
||||
# Translate(destName, srcDS, ...)
|
||||
src_ds = translate_args[1]
|
||||
self.assertEqual(src_ds, mock_mem_ds)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user