geebeam 0.3.2__tar.gz → 0.3.3__tar.gz

{geebeam-0.3.2 → geebeam-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geebeam
-Version: 0.3.2
+Version: 0.3.3
 Summary: Earth Engine + Apache Beam
 License-Expression: MIT
 License-File: LICENSE
@@ -12,6 +12,7 @@ Requires-Dist: earthengine-api>=1.7.10
 Requires-Dist: geopandas>=1.0.1
 Requires-Dist: grpcio>=1.62.3
 Requires-Dist: rasterio>=1.3.10
+Requires-Dist: webdataset>=1.0.2
 Provides-Extra: all
 Requires-Dist: importlib-resources>=6.5.2; extra == 'all'
 Requires-Dist: pytest; extra == 'all'
@@ -75,49 +76,51 @@ PROJECT_ID = google.auth.default()[1]
 # Initialize ee client, replace with your GCP project ID
 ee.Initialize(project=PROJECT_ID)
 
-# Build image for download
-burned_2024 = (ee.ImageCollection('MODIS/061/MCD64A1')
-            .select('BurnDate')
-            .filter(ee.Filter.calendarRange(2024, 2024, 'year'))
-            .min()
-            .gt(0)
-            .rename(['Burn'])
-            )
+### Build image for download
+# Load a raw Landsat 5 ImageCollection for a single year.
+ls5_collection = ee.ImageCollection('LANDSAT/LT05/C02/T1').filterDate(
+    '2010-01-01', '2010-12-31'
+)
+# Create a (mostly) cloud-free Landsat composite
+ls5_composite = ee.Algorithms.Landsat.simpleComposite(
+    ls5_collection,
+    asFloat=True,
+    cloudScoreRange=5)
 
 # Building and triggering the pipeline is done with a single command:
-geebeam.run_pipeline(
-    image_list = [burned_2024],
-    project=PROJECT_ID,
-    patch_size=128, # Pixel dimensions in each direction
-    scale=500, # Final export resolution in meters
+geebeam.sample_and_run_pipeline(
+    image_list = [ls5_composite], # Important: has to be a list of images
+    sampling_region=ee.Geometry.Rectangle(-55.0, -12.0, -50.0, -16.0), # In central-west Brazil
     n_sample=10, # Number of tiles to sample
+    patch_size=128, # Number of pixels in each direction
+    scale=30, # Final export resolution in meters
+    crs='EPSG:4326', # CRS for final output
+    project=PROJECT_ID, # GCP Project ID
+    output_path='./test_data/', # Output path, local or on GCP
     validation_ratio=0.2, # Fraction to select as validation data
-    output_path='./test_tf_data/',
-    sampling_region=ee.Geometry.Rectangle(-63.0, -9.0, -56.0, -4.0)
 )
 ```
 
-Now let's add another dataset: MapBiomas Amazonia forest fraction
+Now let's add another dataset: MapBiomas land-cover from same year. 
+For more info, and legend, see: [MapBiomas Brasil](https://brasil.mapbiomas.org/en/codigos-de-legenda/)
 ```python
-# MB Land-use/land-cover forest fraction
-# Note that LULC codes less than 10 area forest in MapBiomas Amazon Collection 6
-mb_amz_lulc = (
-    ee.Image('projects/mapbiomas-public/assets/amazon/lulc/collection6/mapbiomas_collection60_integration_v1')
-    .lt(10)
-   .reduceResolution('mean', maxPixels=500)
+# MB Land-use/land-cover
+mb_lulc = (
+    ee.Image('projects/mapbiomas-public/assets/brazil/lulc/collection10_1/mapbiomas_brazil_collection10_1_coverage_v1')
+    .select('classification_2010')
 )
 
 # Exporting both together is as simple as this:
-geebeam.run_pipeline(
-    image_list = [burned_2024, mb_amz_lulc],
+geebeam.sample_and_run_pipeline(
+    image_list = [ls5_composite, mb_lulc],
     project=PROJECT_ID,
+    crs='EPSG:4326',
     patch_size=128,
-    scale=500,
+    scale=30,
     n_sample=10,
     validation_ratio=0.2,
-    output_path='./test_tf_data/',
-    sampling_region=ee.Geometry.Rectangle(-63.0, -9.0, -56.0, -4.0),
-    num_workers=1
+    output_path='./test_data_w_mb/',
+    sampling_region=ee.Geometry.Rectangle(-55.0, -12.0, -50.0, -16.0)
 )
 ```
 

{geebeam-0.3.2 → geebeam-0.3.3}/README.md

@@ -37,49 +37,51 @@ PROJECT_ID = google.auth.default()[1]
 # Initialize ee client, replace with your GCP project ID
 ee.Initialize(project=PROJECT_ID)
 
-# Build image for download
-burned_2024 = (ee.ImageCollection('MODIS/061/MCD64A1')
-            .select('BurnDate')
-            .filter(ee.Filter.calendarRange(2024, 2024, 'year'))
-            .min()
-            .gt(0)
-            .rename(['Burn'])
-            )
+### Build image for download
+# Load a raw Landsat 5 ImageCollection for a single year.
+ls5_collection = ee.ImageCollection('LANDSAT/LT05/C02/T1').filterDate(
+    '2010-01-01', '2010-12-31'
+)
+# Create a (mostly) cloud-free Landsat composite
+ls5_composite = ee.Algorithms.Landsat.simpleComposite(
+    ls5_collection,
+    asFloat=True,
+    cloudScoreRange=5)
 
 # Building and triggering the pipeline is done with a single command:
-geebeam.run_pipeline(
-    image_list = [burned_2024],
-    project=PROJECT_ID,
-    patch_size=128, # Pixel dimensions in each direction
-    scale=500, # Final export resolution in meters
+geebeam.sample_and_run_pipeline(
+    image_list = [ls5_composite], # Important: has to be a list of images
+    sampling_region=ee.Geometry.Rectangle(-55.0, -12.0, -50.0, -16.0), # In central-west Brazil
     n_sample=10, # Number of tiles to sample
+    patch_size=128, # Number of pixels in each direction
+    scale=30, # Final export resolution in meters
+    crs='EPSG:4326', # CRS for final output
+    project=PROJECT_ID, # GCP Project ID
+    output_path='./test_data/', # Output path, local or on GCP
     validation_ratio=0.2, # Fraction to select as validation data
-    output_path='./test_tf_data/',
-    sampling_region=ee.Geometry.Rectangle(-63.0, -9.0, -56.0, -4.0)
 )
 ```
 
-Now let's add another dataset: MapBiomas Amazonia forest fraction
+Now let's add another dataset: MapBiomas land-cover from same year. 
+For more info, and legend, see: [MapBiomas Brasil](https://brasil.mapbiomas.org/en/codigos-de-legenda/)
 ```python
-# MB Land-use/land-cover forest fraction
-# Note that LULC codes less than 10 area forest in MapBiomas Amazon Collection 6
-mb_amz_lulc = (
-    ee.Image('projects/mapbiomas-public/assets/amazon/lulc/collection6/mapbiomas_collection60_integration_v1')
-    .lt(10)
-   .reduceResolution('mean', maxPixels=500)
+# MB Land-use/land-cover
+mb_lulc = (
+    ee.Image('projects/mapbiomas-public/assets/brazil/lulc/collection10_1/mapbiomas_brazil_collection10_1_coverage_v1')
+    .select('classification_2010')
 )
 
 # Exporting both together is as simple as this:
-geebeam.run_pipeline(
-    image_list = [burned_2024, mb_amz_lulc],
+geebeam.sample_and_run_pipeline(
+    image_list = [ls5_composite, mb_lulc],
     project=PROJECT_ID,
+    crs='EPSG:4326',
     patch_size=128,
-    scale=500,
+    scale=30,
     n_sample=10,
     validation_ratio=0.2,
-    output_path='./test_tf_data/',
-    sampling_region=ee.Geometry.Rectangle(-63.0, -9.0, -56.0, -4.0),
-    num_workers=1
+    output_path='./test_data_w_mb/',
+    sampling_region=ee.Geometry.Rectangle(-55.0, -12.0, -50.0, -16.0)
 )
 ```
 

{geebeam-0.3.2 → geebeam-0.3.3}/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.3.2'
-__version_tuple__ = version_tuple = (0, 3, 2)
+__version__ = version = '0.3.3'
+__version_tuple__ = version_tuple = (0, 3, 3)
 
 __commit_id__ = commit_id = None

{geebeam-0.3.2 → geebeam-0.3.3}/pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "geopandas>=1.0.1",
     "grpcio>=1.62.3",
     "rasterio>=1.3.10",
+    "webdataset>=1.0.2",
 ]
 
 [project.optional-dependencies]
@@ -57,3 +58,8 @@ include = [
   "pyproject.toml",
   "setup.py"
 ]
+
+[tool.pytest.ini_options]
+markers = [
+    "integration: marks tests as integration tests (slower)",
+]

geebeam-0.3.3/src/geebeam/__init__.py

@@ -0,0 +1,13 @@
+"""Beam and Earth Engine helpers for running data pipelines"""
+
+from . import climate_indices, sampler
+
+from ._pipeline import run_pipeline, sample_and_run_pipeline, grid_and_run_pipeline
+
+__all__ = [
+    "climate_indices",
+    "sampler",
+    "run_pipeline",
+    "sample_and_run_pipeline",
+    "grid_and_run_pipeline"
+]

geebeam-0.3.2/src/geebeam/ee_utils.py → geebeam-0.3.3/src/geebeam/_ee_utils.py

@@ -7,11 +7,11 @@ import io
 import ee
 import numpy as np
 
-def deserialize(obj_json):
+def _deserialize(obj_json):
     """Deserialize Earth Engine JSON DAG"""
     return ee.deserializer.fromJSON(obj_json)
 
-def serialize(obj_ee):
+def _serialize(obj_ee):
     """Serialize Earth Engine object to JSON for Dataflow workers"""
     return ee.serializer.toJSON(obj_ee)
 
@@ -39,7 +39,7 @@ def build_prepped_image(input_list, split_processing=False):
     return prepped_im, band_groups, band_names_flat.getInfo()
 
 def list_to_im(input_list):
-    prepped_im, _ = build_prepped_image(input_list)
+    prepped_im, *_ = build_prepped_image(input_list)
     return prepped_im
 
 def get_pixels(im, point, patch_size, scale_x, scale_y, crs_code):

geebeam-0.3.3/src/geebeam/_pipeline.py

@@ -0,0 +1,334 @@
+"""Prepare and run Beam pipeline to download image 'chips' from Earth Engine"""
+
+import warnings
+import ee
+import geopandas as gpd
+import pandas as pd
+from apache_beam.options.pipeline_options import PipelineOptions
+import numpy as np
+
+from geebeam import _ee_utils, sampler, _transforms
+
+
+def _check_if_localrunner(pipeline_options):
+    """Fixes gRPC timeout issue for local runners."""
+    runner = pipeline_options.get_all_options()['runner']
+    if runner is None or runner in ['DirectRunner', 'PrismRunner']:
+        return True
+    else:
+        return False
+
+def _type_inference(val):
+    if isinstance(val, int):
+        return 'int'
+    elif isinstance(val, float):
+        return 'float'
+    elif isinstance(val, list) or isinstance(val, np.ndarray):
+        return {'arraylike': np.array(val).shape}
+    elif isinstance(val, str):
+        return 'str'
+    else:
+        raise ValueError
+
+def _build_md_feature_dict(record, extra_metadata):
+    md_feature_dict = {
+        'id': 'int',
+        'x': 'float',
+        'y': 'float',
+        'split':'str'
+    }
+    if extra_metadata is not None:
+        for key in extra_metadata.keys():
+            # Basic type inference for extra metadata
+            try:
+                md_type = _type_inference(extra_metadata[key])
+            except ValueError:
+                raise ValueError(f'Could not determine data type of extra_metadata feature {key}')
+            md_feature_dict[key] = md_type
+
+    for key in record.keys():
+        if key not in ['id','x','y','split']:
+            try:
+                md_type = _type_inference(record[key])
+            except ValueError:
+                raise ValueError(f'Could not determine data type of column feature {key}')
+            md_feature_dict[key] = md_type
+    return md_feature_dict
+
+def _prepare_run_metadata(config):
+    ee.Initialize(project=config['project_id'])
+
+    proj = ee.Projection(config['crs']).atScale(config['scale'])
+    proj_dict = proj.getInfo()
+
+    scale_x = proj_dict['transform'][0]
+    scale_y = -proj_dict['transform'][4]
+
+    return scale_x, scale_y
+
+def run_pipeline(
+        image_list: list[ee.Image],
+        output_path: str,
+        project: str,
+        patch_size: int,
+        scale: float,
+        sampling_points: pd.DataFrame | gpd.GeoDataFrame | ee.FeatureCollection,
+        output_type: str = 'tiff',
+        crs: str = 'EPSG:4326',
+        split_processing: bool = False,
+        extra_metadata: dict = {},
+        beam_options: dict[str] | list[str] | None = None,
+        dataset_version: str = '1.0.0',
+        dataset_name: str = 'geebeam_dataset'
+        ) -> None:
+    """Run a Beam pipeline to download image chips from Earth Engine.
+
+    Args:
+        image_list: A list of image identifiers to process.
+        sampling_points: Locations to sample from, specifying upper-left of box.
+        output_path: The path where output will be saved.
+        output_type: 'tiff' (tiffs with parquet for metadata),
+            'webdataset' (tiffs with jsons, in sharded tars),
+            'tfrecord' (raw tfrecords), or 'tfds' (tensorflow-dataset).
+        project: The Google Cloud project ID.
+        patch_size: The size of the patches to be processed.
+        scale: The scale factor for image processing.
+        crs: The coordinate reference system. Defaults to 'EPSG:4326'.
+        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.
+    """
+    import logging
+
+    logging.getLogger().setLevel(logging.INFO)
+
+    # Set up configuration dict to pass along
+    config = {
+        'project_id': project,
+        'patch_size': patch_size,
+        'scale': scale,
+        'crs': crs
+    }
+
+    # Parses from command line and/or retrieves from dict. Note that dict takes precedent.
+    if isinstance(beam_options, dict):
+        pipeline_options = PipelineOptions(
+            **beam_options,
+            project=config['project_id'],
+            save_main_session=True,
+            )
+    elif isinstance(beam_options, list):
+        warnings.warn('Creating PipelineOptions from beam_options list.'
+                      ' Ignores command-line beam options.')
+        pipeline_options = PipelineOptions(
+            beam_options,
+            project=config['project_id'],
+            save_main_session=True,
+            )
+    else:
+        pipeline_options = PipelineOptions(
+            project=config['project_id'],
+            save_main_session=True,
+            )
+
+    # Get sample points
+    input_records, splits = sampler._process_sampling_points(sampling_points, target_crs=config['crs'])
+
+    # Get types of extra non-image data in extra_medtada or input_records
+    md_feature_dict = _build_md_feature_dict(input_records[0], extra_metadata)
+
+    # Pre-run info:
+    scale_x, scale_y = _prepare_run_metadata(config)
+
+    # Check if a local runner. If so, add longer job timeout to fix grpcio timeout issue
+    is_local = _check_if_localrunner(pipeline_options)
+    if is_local:
+        logging.warning('Running on local runner. Setting beam job_server_timeout'
+                        ' to 9999999 seconds to avoid grpcio timeout errors.')
+        pipeline_options_dict = pipeline_options.get_all_options(drop_default=True)
+        pipeline_options_dict['job_server_timeout'] = 9999999
+        pipeline_options = PipelineOptions.from_dictionary(pipeline_options_dict)
+
+    # Prepare and serialize inputs
+    # band_groups is a list of lists containing bands to export
+    # if split_processing = False, will contain one list with all bands
+    # if split_processing = True,  contains separate band_lists for each image in image_list
+    prepped_image, band_groups, all_bands = _ee_utils.build_prepped_image(image_list, split_processing=split_processing)
+    serialized_image = _ee_utils._serialize(prepped_image)
+
+    # Execute pipeline based on output type:
+    if output_type == 'tfrecord':
+        try:
+            from geebeam import _tfrecord_writer
+        except ImportError:
+            raise ImportError(
+                "Missing dependencies for tfrecord writer. "
+                "Install them with `pip install geebeam[tensorflow]`"
+            )
+        # Write sidecar schema before pipeline execution
+        extra_keys = list(extra_metadata.keys())
+        _transforms._write_sidecar_schema(output_path, all_bands, extra_keys,
+                                        is_gcs=output_path.startswith('gs://'))
+        _tfrecord_writer.run_tfrecord_export(
+            input_records=input_records,
+            splits=splits,
+            output_path=output_path,
+            config=config,
+            serialized_image=serialized_image,
+            band_groups=band_groups,
+            scale_x=scale_x,
+            scale_y=scale_y,
+            extra_metadata=extra_metadata,
+            md_feature_dict=md_feature_dict,
+            pipeline_options=pipeline_options
+        )
+    elif output_type == 'tfds':
+        try:
+            from geebeam import _tfds_writer
+        except ImportError:
+            raise ImportError(
+                "Missing dependencies for TFDS writer. "
+                "Install them with `pip install geebeam[tensorflow]`"
+            )
+        _tfds_writer.run_tfds_export(
+            input_records=input_records,
+            splits=splits,
+            output_path=output_path,
+            config=config,
+            serialized_image=serialized_image,
+            band_groups=band_groups,
+            all_bands=all_bands,
+            scale_x=scale_x,
+            scale_y=scale_y,
+            extra_metadata=extra_metadata,
+            md_feature_dict=md_feature_dict,
+            pipeline_options=pipeline_options,
+            dataset_name=dataset_name,
+            dataset_version=dataset_version
+        )
+    elif output_type == 'tif' or output_type == 'tiff':
+        from geebeam import _tiff_writer
+        _tiff_writer.run_tiff_export(
+            input_records=input_records,
+            splits=splits,
+            output_path=output_path,
+            config=config,
+            serialized_image=serialized_image,
+            band_groups=band_groups,
+            scale_x=scale_x,
+            scale_y=scale_y,
+            extra_metadata=extra_metadata,
+            md_feature_dict=md_feature_dict,
+            pipeline_options=pipeline_options
+        )
+    elif output_type == 'webdataset':
+        from geebeam import _wds_writer
+        _wds_writer.run_webdataset_export(
+            input_records=input_records,
+            splits=splits,
+            output_path=output_path,
+            config=config,
+            serialized_image=serialized_image,
+            band_groups=band_groups,
+            scale_x=scale_x,
+            scale_y=scale_y,
+            extra_metadata=extra_metadata,
+            pipeline_options=pipeline_options
+        )
+    else:
+        raise ValueError(f"output_type {output_type} not implemented. Options are ['tfds', 'tfrecord', 'tiff', 'webdataset']")
+
+def sample_and_run_pipeline(
+        sampling_region: str | gpd.GeoDataFrame | ee.Geometry,
+        n_sample: int,
+        validation_ratio: float = 0,
+        random_seed: int = 0,
+        crs: str = 'EPSG:4326',
+        *args,
+        **kwargs
+        ) -> None:
+    """Sample random points and then run a Beam pipeline to download image chips from Earth Engine.
+
+    Args:
+        image_list: A list of image identifiers to process.
+        sampling_region: Region to sample from, polygon or group of polygons.
+        n_sample: Number of points to sample.
+        validation_ratio: Fraction of points to mark as validation.
+        random_seed: Seed for random sampling
+        output_path: The path where output will be saved.
+        output_type: 'tiff' (tiffs with parquet for metadata),
+            'webdataset' (tiffs with jsons, in sharded tars),
+            'tfrecord' (raw tfrecords), or 'tfds' (tensorflow-dataset).
+        project: The Google Cloud project ID.
+        patch_size: The size of the patches to be processed.
+        scale: The scale factor for image processing.
+        crs: The coordinate reference system. Defaults to 'EPSG:4326'.
+        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.
+    """
+
+    sample_points = sampler.sample_region_random(
+        roi=sampling_region,
+        n_sample=n_sample,
+        random_seed=random_seed,
+        crs=crs,
+    )
+
+    if validation_ratio > 0:
+        sample_points = sampler.split_sets(
+            sample_points, split_names=['train','validation'], split_ratios=[1-validation_ratio,validation_ratio],
+            random_seed=random_seed
+        )
+
+    return run_pipeline(sampling_points=sample_points, crs=crs, *args, **kwargs)
+
+def grid_and_run_pipeline(
+        sampling_region: str | gpd.GeoDataFrame | ee.Geometry,
+        validation_ratio: float,
+        scale: float,
+        stride: int,
+        buffer_distance: float = 0,
+        crs: str = 'EPSG:4326',
+        random_seed: int = 0,
+        *args,
+        **kwargs
+        ) -> None:
+    """Sample points from regular grid and then run a Beam pipeline to download image chips from Earth Engine.
+
+    Args:
+        image_list: A list of image identifiers to process.
+        sampling_region: Region to sample from, polygon or group of polygons.
+        validation_ratio: Fraction of points to mark as validation.
+        stride: Number of pixels between consecutive samples. If want full coverage without overlaps,
+            stride should be equal to patch_size. If less than patch_size, will generate overlaps.
+            If greater, will be gaps between sampled patches.
+        random_seed: Seed for random sampling
+        output_path: The path where output will be saved.
+        output_type: 'tiff' (tiffs with parquet for metadata),
+            'webdataset' (tiffs with jsons, in sharded tars),
+            'tfrecord' (raw tfrecords), or 'tfds' (tensorflow-dataset).
+        project: The Google Cloud project ID.
+        patch_size: The size of the patches to be processed.
+        scale: The scale factor for image processing.
+        crs: The coordinate reference system. Defaults to 'EPSG:4326'.
+        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.
+    """
+
+    sample_points = sampler.sample_region_grid(
+        roi=sampling_region,
+        crs='EPSG:4326',
+        stride=stride,
+        scale=scale,
+        buffer_distance=buffer_distance,
+    )
+
+    if validation_ratio > 0:
+        sample_points = sampler.split_sets(
+            sample_points, split_names=['train','validation'], split_ratios=[1-validation_ratio,validation_ratio],
+            random_seed=random_seed
+        )
+
+    return run_pipeline(sampling_points=sample_points, crs=crs, scale=scale, *args, **kwargs)

geebeam-0.3.2/src/geebeam/tf_utils.py → geebeam-0.3.3/src/geebeam/_tf_utils.py

@@ -2,8 +2,9 @@
 
 import tensorflow as tf
 import apache_beam as beam
+import numpy as np
 
-from geebeam import transforms
+from geebeam import _transforms
 
 def _bytes_feature(value):
     """Returns a bytes_list from a string / byte."""
@@ -17,28 +18,39 @@ def _int64_feature(value):
     """Returns an int64_list from a bool / enum / int / uint."""
     return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
 
-def dict_to_example(element):
+def _dict_to_example(record):
     """"Convert structured numpy array to tf.Example proto."""
     # First add metadata
-    md_dict = {
-        'md_id': _int64_feature(element['metadata']['id']),
-        'md_y': _float_feature(element['metadata']['y']),
-        'md_x': _float_feature(element['metadata']['x']),
-        }
-    for md_key in element['metadata'].keys():
-        if md_key not in ['id','y','x', 'split']:
-            md_dict['md_' + md_key] = tf.train.Feature(float_list=
+    # Add metadata features
+    md_dict = {}
+    for md_key, md_val in record['metadata'].items():
+        if isinstance(md_val, str):
+            md_dict[f'md_{md_key}'] = _bytes_feature(md_val.encode('utf-8'))
+        elif np.isscalar(md_val):
+            if isinstance(md_val, int):
+                md_dict[f'md_{md_key}'] = _int64_feature(record['metadata'][md_key])
+            else:
+                md_dict[f'md_{md_key}'] = _float_feature(record['metadata'][md_key])
+        else:
+            md_dict[f'md_{md_key}'] = tf.train.Feature(float_list=
                 tf.train.FloatList(
-                    value = transforms.convert_to_iterable(element['metadata'][md_key])
+                    value = _transforms._convert_to_iterable(record['metadata'][md_key])
                 )
             )
 
     # Build image feature with named bands
     array_dict = {}
-    for im_feat in element['array'].keys():#.dtype.names:
+    for band_name, band_data in record['array'].items():
+        array_dict[f'{band_name}'] = band_data.flatten().astype('float32')
+
+    # Build image feature with named bands
+    array_dict = {}
+    for im_feat, im_data in record['array'].items():
         array_dict['im_'+im_feat] = tf.train.Feature(
             float_list = tf.train.FloatList(
-                value = element['array'][im_feat].flatten()))
+                value = im_data.flatten()
+                )
+        )
 
     # Combine
     feature = {**md_dict, **array_dict}
@@ -47,16 +59,6 @@ def dict_to_example(element):
     return tf.train.Example(
         features = tf.train.Features(feature = feature)).SerializeToString()
 
-def array_to_example(structured_array):
-    """"Convert structured numpy array to tf.Example proto."""
-    feature = {}
-    for f in structured_array.dtype.names:
-        feature[f] = tf.train.Feature(
-            float_list = tf.train.FloatList(
-                value = structured_array[f].flatten()))
-    return tf.train.Example(
-        features = tf.train.Features(feature = feature))
-
 class WriteTFExample(beam.PTransform):
     """Write example"""
     def __init__(self, output_dir, file_name_suffix='.tfrecord.gz'):