Tiff dtype

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']
                 ))
             )
 

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))
             )

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.

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']]

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'):

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)

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