From 6c8c684c3c6bdafaca8bebbddccddc5b426d2a37 Mon Sep 17 00:00:00 2001 From: Kylen Solvik Date: Fri, 24 Jul 2026 11:41:35 -0400 Subject: [PATCH 1/2] Manually specify dtype for tiff and wds writer (defaults to float32) --- src/geebeam/_tiff_writer.py | 13 ++++++++----- src/geebeam/_wds_writer.py | 20 ++++++++++++-------- src/geebeam/pipeline.py | 13 +++++++++++-- 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/src/geebeam/_tiff_writer.py b/src/geebeam/_tiff_writer.py index b737ba2..5bf5f9e 100644 --- a/src/geebeam/_tiff_writer.py +++ b/src/geebeam/_tiff_writer.py @@ -5,6 +5,7 @@ import tempfile import apache_beam as beam +import numpy as np import pyarrow as pa import rasterio from apache_beam.io.filesystems import FileSystems @@ -21,11 +22,12 @@ def _build_tiff_name(id, min_digits=5): class WriteTiff(beam.DoFn): """DoFn to write image arrays to Cloud-Optimized GeoTIFFs.""" - def __init__(self, output_path, crs, scale_x, scale_y): + def __init__(self, output_path, crs, scale_x, scale_y, output_dtype='float32'): self.output_path = output_path self.crs = crs self.scale_x = scale_x self.scale_y = scale_y + self.dtype = np.dtype(output_dtype) def setup(self): # Ensure the output directory for TIFFs exists @@ -47,8 +49,8 @@ def process(self, element): first_band = next(iter(array_dict.values())) height, width = first_band.shape count = len(array_dict) - dtype = first_band.dtype - + dtype = self.dtype + # Construct affine transform transform = Affine( self.scale_x, 0, metadata.get('x_topleft', metadata['x']), @@ -71,7 +73,7 @@ def process(self, element): compress='lzw' ) as dst: for i, (band_name, data) in enumerate(array_dict.items(), 1): - dst.write(data, i) + dst.write(data.astype(dtype), i) dst.set_band_description(i, band_name) # Upload the temporary file to the final destination @@ -152,7 +154,8 @@ def run_tiff_export( output_path=output_dir, crs=config['crs'], scale_x=scale_x, - scale_y=scale_y + scale_y=scale_y, + output_dtype=config['output_dtype'] )) ) diff --git a/src/geebeam/_wds_writer.py b/src/geebeam/_wds_writer.py index d2d5485..e4bbc00 100644 --- a/src/geebeam/_wds_writer.py +++ b/src/geebeam/_wds_writer.py @@ -5,6 +5,7 @@ import uuid import apache_beam as beam +import numpy as np import webdataset as wds from apache_beam.options.pipeline_options import PipelineOptions from rasterio.io import MemoryFile @@ -13,13 +14,13 @@ from geebeam import _transforms -def _create_tiff_bytes(array_dict, metadata, crs, scale_x, scale_y): +def _create_tiff_bytes(array_dict, metadata, crs, scale_x, scale_y, output_dtype='float32'): """Create TIFF bytes from array dict and metadata.""" first_band = next(iter(array_dict.values())) height, width = first_band.shape count = len(array_dict) - dtype = first_band.dtype - + dtype = np.dtype(output_dtype) + transform = Affine( scale_x, 0, metadata.get('x_topleft', metadata['x']), 0, scale_y, metadata.get('y_topleft', metadata['y']) @@ -38,23 +39,25 @@ def _create_tiff_bytes(array_dict, metadata, crs, scale_x, scale_y): compress='lzw' ) as dst: for i, (band_name, data) in enumerate(array_dict.items(), 1): - dst.write(data, i) + dst.write(data.astype(dtype), i) dst.set_band_description(i, band_name) return memfile.read() class ProcessToWebDataset(beam.DoFn): """DoFn to prepare records for WebDataset output.""" - def __init__(self, crs, scale_x, scale_y): + def __init__(self, crs, scale_x, scale_y, output_dtype='float32'): self.crs = crs self.scale_x = scale_x self.scale_y = scale_y + self.output_dtype = output_dtype def process(self, element): metadata = element['metadata'] array_dict = element['array'] basename = str(metadata['id']).zfill(5) - - tif_bytes = _create_tiff_bytes(array_dict, metadata, self.crs, self.scale_x, self.scale_y) + + tif_bytes = _create_tiff_bytes(array_dict, metadata, self.crs, self.scale_x, self.scale_y, + output_dtype=self.output_dtype) json_bytes = json.dumps(metadata).encode('utf-8') @@ -118,7 +121,8 @@ def run_webdataset_export( | f'Format {split}' >> beam.ParDo(ProcessToWebDataset( crs=config['crs'], scale_x=scale_x, - scale_y=scale_y + scale_y=scale_y, + output_dtype=config['output_dtype'] )) | f'Write {split}' >> beam.ParDo(WriteToWebDataset(output_path, split)) ) diff --git a/src/geebeam/pipeline.py b/src/geebeam/pipeline.py index bd1a144..e8fb41e 100644 --- a/src/geebeam/pipeline.py +++ b/src/geebeam/pipeline.py @@ -114,6 +114,7 @@ def run_pipeline( crs: str = 'EPSG:4326', align_transform: Affine | tuple[float] | list[float] | None = None, output_type: str = 'tiff', + output_dtype: str = 'float32', split_processing: bool = False, extra_metadata: dict | None = None, beam_options: dict[str] | list[str] | None = None, @@ -143,6 +144,9 @@ def run_pipeline( output_type: 'tiff' (tiffs with parquet for metadata), 'webdataset' (tiffs with jsons, in sharded tars), 'tfrecord' (raw tfrecords), or 'tfds' (tensorflow-dataset). + output_dtype: dtype for GeoTIFF outputs ('tiff'/'webdataset'). A GeoTIFF uses a single + dtype for all bands, so every band is cast to this type. Defaults to 'float32'. + Ignored for 'tfrecord'/'tfds' (which store bands as float lists). split_processing: Flag to indicate if processing should be split. Defaults to False. extra_metadata: Additional metadata to include. Defaults to an empty dictionary. beam_options_dict: Options for the Beam pipeline. Defaults to an empty dictionary. @@ -156,7 +160,8 @@ def run_pipeline( """ import logging - logger = logging.getLogger(__name__).setLevel(logging.INFO) + logger = logging.getLogger(__name__) + logger.setLevel(logging.INFO) if isinstance(image_list, ee.ImageCollection): raise TypeError( @@ -188,12 +193,16 @@ def run_pipeline( if not extra_metadata: extra_metadata = {} + # Validate output_dtype early (raises TypeError for an unrecognized dtype string) + np.dtype(output_dtype) + # Set up configuration dict to pass along config = { 'project_id': project, 'patch_size': patch_size, 'scale': scale, - 'crs': crs + 'crs': crs, + 'output_dtype': output_dtype } # Parses from command line and/or retrieves from dict. Note that dict takes precedent. From eff5e246375b98d5d9f56fa122962aa3ad61a2cf Mon Sep 17 00:00:00 2001 From: Kylen Solvik Date: Fri, 24 Jul 2026 11:44:04 -0400 Subject: [PATCH 2/2] Updated tests with output_dtype flag --- tests/test_integration.py | 2 ++ tests/test_pipeline.py | 13 +++++++++++++ tests/test_tiff_writer.py | 32 ++++++++++++++++++++++++++++++++ tests/test_wds_writer.py | 18 ++++++++++++++++++ 4 files changed, 65 insertions(+) diff --git a/tests/test_integration.py b/tests/test_integration.py index d381072..496a141 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -49,6 +49,7 @@ def test_tiff_pipeline_integration(tmp_path): 'patch_size': 4, 'scale': 10.0, 'crs': 'EPSG:4326', + 'output_dtype': 'float32' } serialized_image = '{"type": "Image"}' band_groups = [['band1']] @@ -97,6 +98,7 @@ def test_webdataset_pipeline_integration(tmp_path): 'patch_size': 4, 'scale': 10.0, 'crs': 'EPSG:4326', + 'output_dtype': 'float32' } serialized_image = '{"type": "Image"}' band_groups = [['band1']] diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index fea6b32..b9701b8 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -202,6 +202,19 @@ def test_run_pipeline_requires_scale_or_align(): sampling_points=MagicMock(), ) +def test_run_pipeline_invalid_output_dtype(): + """An unrecognized output_dtype should fail fast via np.dtype.""" + with pytest.raises(TypeError): + run_pipeline( + image_list=[MagicMock()], + output_path='/tmp/test', + project='test-project', + patch_size=4, + scale=30.0, + sampling_points=MagicMock(), + output_dtype='not-a-real-dtype', + ) + def test_run_pipeline_rejects_image_collection(): """An ee.ImageCollection passed as image_list should raise TypeError.""" with pytest.raises(TypeError, match='ee.ImageCollection'): diff --git a/tests/test_tiff_writer.py b/tests/test_tiff_writer.py index 8921eac..ec92b60 100644 --- a/tests/test_tiff_writer.py +++ b/tests/test_tiff_writer.py @@ -57,6 +57,38 @@ def test_write_tiff_process_fallback_to_xy(tmp_path): assert ds.transform.c == pytest.approx(10.0) assert ds.transform.f == pytest.approx(20.0) +def test_write_tiff_heterogeneous_dtypes_not_truncated(tmp_path): + """Bands of different native dtypes are cast to output_dtype (default float32) without loss.""" + writer = WriteTiff(output_path=str(tmp_path), crs='EPSG:4326', scale_x=0.0001, scale_y=-0.0001) + writer.setup() + + element = { + 'metadata': {'id': 8, 'x': 10.0, 'y': 20.0, 'split': 'train'}, + 'array': { + 'mask_uint8': np.ones((4, 4), dtype=np.uint8), + 'frac_float64': np.full((4, 4), 0.5, dtype=np.float64), + }, + } + writer.process(element) + + with rasterio.open(os.path.join(str(tmp_path), '00008.tif')) as ds: + assert ds.dtypes == ('float32', 'float32') + assert ds.read(2).min() == 0.5 # float band survives (would be 0 if cast to uint8) + +def test_write_tiff_respects_output_dtype(tmp_path): + """output_dtype is honored (e.g. float64).""" + writer = WriteTiff(output_path=str(tmp_path), crs='EPSG:4326', scale_x=0.0001, scale_y=-0.0001, + output_dtype='float64') + writer.setup() + element = { + 'metadata': {'id': 9, 'x': 10.0, 'y': 20.0, 'split': 'train'}, + 'array': {'band1': np.ones((4, 4), dtype=np.uint8)}, + } + writer.process(element) + + with rasterio.open(os.path.join(str(tmp_path), '00009.tif')) as ds: + assert ds.dtypes == ('float64',) + def test_process_metadata_to_parquet(tmp_path): output_dir = str(tmp_path) dofn = ProcessMetadataToParquet(output_path=output_dir) diff --git a/tests/test_wds_writer.py b/tests/test_wds_writer.py index 80f7197..cb71f8a 100644 --- a/tests/test_wds_writer.py +++ b/tests/test_wds_writer.py @@ -26,6 +26,24 @@ def test_create_tiff_bytes(): assert ds.width == 4 assert ds.height == 4 +def test_create_tiff_bytes_heterogeneous_dtypes_not_truncated(): + """Bands with different native dtypes must be cast to output_dtype without truncation + (regression: the writer used to force every band to the first band's dtype).""" + array_dict = { + 'mask_uint8': np.ones((4, 4), dtype=np.uint8), # first band -> used to win + 'frac_float64': np.full((4, 4), 0.5, dtype=np.float64), # would truncate to 0 as uint8 + } + metadata = {'id': 0, 'x': 10.0, 'y': 20.0, 'split': 'train'} + + result = _create_tiff_bytes(array_dict, metadata, 'EPSG:4326', 0.0001, -0.0001, + output_dtype='float32') + + with rasterio.open(io.BytesIO(result)) as ds: + assert ds.count == 2 + assert set(ds.dtypes) == {'float32'} + assert ds.read(1).max() == 1.0 # uint8 mask preserved + assert ds.read(2).min() == 0.5 # float band NOT truncated to 0 + def test_process_to_webdataset(): crs = 'EPSG:4326' scale_x = 0.0001