geebeam 0.2.2__tar.gz → 0.3.2__tar.gz

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

@@ -1,19 +1,39 @@
 Metadata-Version: 2.4
 Name: geebeam
-Version: 0.2.2
+Version: 0.3.2
 Summary: Earth Engine + Apache Beam
 License-Expression: MIT
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: apache-beam[gcp]>=2.70.0
+Requires-Dist: apache-beam[gcp]==2.72.0
+Requires-Dist: crcmod>=1.7
 Requires-Dist: dill>=0.4.1
 Requires-Dist: earthengine-api>=1.7.10
 Requires-Dist: geopandas>=1.0.1
 Requires-Dist: grpcio>=1.62.3
-Requires-Dist: tensorflow-data-validation>=1.16
-Requires-Dist: tensorflow-metadata>=1.16
-Requires-Dist: tensorflow>=2.16
-Requires-Dist: tfx-bsl>=1.16
+Requires-Dist: rasterio>=1.3.10
+Provides-Extra: all
+Requires-Dist: importlib-resources>=6.5.2; extra == 'all'
+Requires-Dist: pytest; extra == 'all'
+Requires-Dist: ruff; extra == 'all'
+Requires-Dist: tensorflow-data-validation>=1.16; extra == 'all'
+Requires-Dist: tensorflow-datasets>=4.9.9; extra == 'all'
+Requires-Dist: tensorflow-metadata>=1.16; extra == 'all'
+Requires-Dist: tensorflow>=2.16; extra == 'all'
+Requires-Dist: tfx-bsl>=1.16; extra == 'all'
+Provides-Extra: build
+Requires-Dist: hatch; extra == 'build'
+Provides-Extra: lint
+Requires-Dist: ruff; extra == 'lint'
+Provides-Extra: tensorflow
+Requires-Dist: importlib-resources>=6.5.2; extra == 'tensorflow'
+Requires-Dist: tensorflow-data-validation>=1.16; extra == 'tensorflow'
+Requires-Dist: tensorflow-datasets>=4.9.9; extra == 'tensorflow'
+Requires-Dist: tensorflow-metadata>=1.16; extra == 'tensorflow'
+Requires-Dist: tensorflow>=2.16; extra == 'tensorflow'
+Requires-Dist: tfx-bsl>=1.16; extra == 'tensorflow'
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
 Description-Content-Type: text/markdown
 
 # GeeBeam
@@ -114,7 +134,7 @@ python examples/geebeam_run.py \
     --experiments=use_runner_v2 \
     --temp_location=gs://[your-bucket]/[path_to_temp_dir]
     --machine_type=n2-highmem-2 \
-    --sdk_container_image=us-docker.pkg.dev/mmacedo-reservoirid/geebeam-public/geebeam:latest
+    --sdk_container_image=kysolvik/geebeam:latest
 ```
 
 Note in this case your `output_path` in run_pipeline() should be a Google Cloud Storage path. If you're running an older version of geebeam, replace "latest" in the sdk_container_image URI with the version number (e.g. v0.1.2). You can also build your own Docker image to run on. More info in the [DataFlow docs](https://docs.cloud.google.com/dataflow/docs/guides/using-custom-containers).

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

@@ -96,7 +96,7 @@ python examples/geebeam_run.py \
     --experiments=use_runner_v2 \
     --temp_location=gs://[your-bucket]/[path_to_temp_dir]
     --machine_type=n2-highmem-2 \
-    --sdk_container_image=us-docker.pkg.dev/mmacedo-reservoirid/geebeam-public/geebeam:latest
+    --sdk_container_image=kysolvik/geebeam:latest
 ```
 
 Note in this case your `output_path` in run_pipeline() should be a Google Cloud Storage path. If you're running an older version of geebeam, replace "latest" in the sdk_container_image URI with the version number (e.g. v0.1.2). You can also build your own Docker image to run on. More info in the [DataFlow docs](https://docs.cloud.google.com/dataflow/docs/guides/using-custom-containers).

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

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

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

@@ -10,15 +10,35 @@ description = "Earth Engine + Apache Beam"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-    "apache-beam[gcp]>=2.70.0",
+    "apache-beam[gcp]==2.72.0",
+    "crcmod>=1.7",
     "dill>=0.4.1",
     "earthengine-api>=1.7.10",
     "geopandas>=1.0.1",
+    "grpcio>=1.62.3",
+    "rasterio>=1.3.10",
+]
+
+[project.optional-dependencies]
+test = ["pytest"]
+lint = ["ruff"]
+build = ["hatch"]
+tensorflow = [
     "tensorflow>=2.16",
     "tensorflow-data-validation>=1.16",
     "tensorflow-metadata>=1.16",
     "tfx-bsl>=1.16",
-    "grpcio>=1.62.3",
+    "tensorflow-datasets>=4.9.9",
+    "importlib-resources>=6.5.2",
+]
+all = [
+    "geebeam[tensorflow,test,lint]"
+]
+
+# Pin upper version of python, since apache-beam only supports to 3.12
+[tool.uv]
+environments = [
+    "python_version < '3.13'"
 ]
 
 # Version control
@@ -28,7 +48,6 @@ version-file = "_version.py"
 [tool.hatch.version]
 source = "vcs"
 
-
 [tool.hatch.build.targets.sdist]
 include = [
   "src",

geebeam-0.3.2/src/geebeam/__init__.py

@@ -0,0 +1,10 @@
+"""Beam and Earth Engine helpers for running data pipelines"""
+
+from . import climate_indices
+
+from .pipeline import run_pipeline
+
+__all__ = [
+    "climate_indices",
+    "run_pipeline",
+]

geebeam-0.3.2/src/geebeam/ee_utils.py

@@ -0,0 +1,94 @@
+"""
+Utilities for cleaning, combining, and serializing/deserializing EE objects.
+"""
+
+import io
+
+import ee
+import numpy as np
+
+def deserialize(obj_json):
+    """Deserialize Earth Engine JSON DAG"""
+    return ee.deserializer.fromJSON(obj_json)
+
+def serialize(obj_ee):
+    """Serialize Earth Engine object to JSON for Dataflow workers"""
+    return ee.serializer.toJSON(obj_ee)
+
+def get_band_names(input_list):
+    """Get simplified band_names for output (without prefixed image_id)
+
+    TODO: Add year to distinguish multiple years?
+    """
+    return [
+        image.bandNames()
+        for image in input_list
+    ]
+
+def build_prepped_image(input_list, split_processing=False):
+    """Combine a list of EE images into single image"""
+    band_names = get_band_names(input_list)
+    band_names_flat = ee.List(band_names).flatten()
+
+    if split_processing:
+        band_groups = band_names
+    else:
+        band_groups = [band_names_flat]
+    # Final prepped image
+    prepped_im = ee.ImageCollection(input_list).toBands().rename(band_names_flat)
+    return prepped_im, band_groups, band_names_flat.getInfo()
+
+def list_to_im(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):
+    # Make a request object.
+    request = {
+        'expression': im,
+        'fileFormat': 'NPY',
+        'grid': {
+            'dimensions': {
+                'width': patch_size,
+                'height': patch_size
+            },
+            'affineTransform': {
+                'scaleX': scale_x,
+                'shearX': 0,
+                'translateX': point['x'],
+                'shearY': 0,
+                'scaleY': scale_y,
+                'translateY': point['y']
+            },
+            'crsCode': crs_code
+        },
+    }
+
+    raw = ee.data.computePixels(request)
+
+    if not raw:
+        raise RuntimeError("Empty EE response")
+
+    try:
+        arr = np.load(io.BytesIO(raw), allow_pickle=True)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load NPY for coords {point['id']}") from e
+
+    return arr
+
+def get_pixels_allbands(im, band_groups, point, patch_size, scale_x, scale_y, crs_code):
+    """Get pixels for all groups of bands"""
+    out_ars = []
+    for band_list in band_groups:
+        prepped_image = im.select(band_list)
+        out_ars.append(
+            get_pixels(
+            im=prepped_image,
+            point=point,
+            patch_size=patch_size,
+            scale_x=scale_x,
+            scale_y=scale_y,
+            crs_code=crs_code
+            )
+        )
+    return out_ars

{geebeam-0.2.2 → geebeam-0.3.2}/src/geebeam/pipeline.py

@@ -5,12 +5,7 @@ import ee
 import geopandas as gpd
 import pandas as pd
 import numpy as np
-import apache_beam as beam
 from apache_beam.options.pipeline_options import PipelineOptions
-import tensorflow_data_validation as tfdv
-from tfx_bsl.coders import example_coder
-import os
-from google.protobuf.json_format import MessageToJson
 
 from geebeam import ee_utils, sampler, transforms
 
@@ -42,6 +37,7 @@ def run_pipeline(
         patch_size: int,
         scale: float,
         n_sample: int,
+        output_type: str = 'tiff',
         sampling_region: str | gpd.GeoDataFrame | ee.Geometry | None = None,
         sampling_points: pd.DataFrame | gpd.GeoDataFrame | ee.Geometry | None = None,
         validation_ratio: float = 0.2,
@@ -49,7 +45,9 @@ def run_pipeline(
         random_seed: int = None,
         split_processing: bool = False,
         extra_metadata: dict = {},
-        beam_options: dict[str] | list[str] | None = None
+        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.
 
@@ -58,6 +56,7 @@ def run_pipeline(
         sampling_region: The region for sampling images.
         sampling_points: Locations to sample from, specifying upper-left of box.
         output_path: The path where output will be saved.
+        output_type: '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.
@@ -121,110 +120,81 @@ def run_pipeline(
     # Pre-run info:
     scale_x, scale_y = _prepare_run_metadata(config)
 
-    # Set up pipeline
-
-    # Check if a local runner
+    # 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 = ee_utils.build_prepped_image(image_list, split_processing=split_processing)
+    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
-    with beam.Pipeline(options=pipeline_options) as pipeline:
-
-        points = pipeline | 'Create points' >> beam.Create(input_records)
 
-        if is_local:
-            batches = (
-                points
-                | 'Add Dummy Key' >> beam.Map(lambda x: (None, x))
-                | 'Reshuffle' >> beam.Reshuffle()
-                | 'Force Single Batches' >> beam.GroupIntoBatches(batch_size=1)
-                | 'Extract' >> beam.FlatMap(lambda x: x[1])
+    # 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]`"
             )
-        else:
-            batches = points
-
-        training_data, validation_data = (
-            batches
-            | 'Get patch' >> beam.ParDo(transforms.EEComputePatch(
-                config,
-                serialized_image,
-                scale_x,
-                scale_y,
-                band_groups
-                ))
-            | 'Add metadata' >> beam.ParDo(transforms.AddMetadata(extra_metadata))
-            | 'Split dataset' >> beam.Partition(transforms.split_dataset, 2)
-        )
-
-        # Convert to TF examples
-        training_examples = (
-            training_data
-            | 'Train to tf.Example' >> beam.Map(transforms.dict_to_example)
+        # 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,
+            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
         )
-        # Calculate stats on training data
-        decoder = example_coder.ExamplesToRecordBatchDecoder()
-        stats = (
-            training_examples
-            | 'Batch' >> beam.BatchElements(
-                min_batch_size=10,
-                max_batch_size=100)
-            | 'Decode to arrow' >> beam.Map(lambda b: decoder.DecodeBatch(b))
-            | 'Generate Statistics' >> tfdv.GenerateStatistics()
+    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,
+            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,
+            pipeline_options=pipeline_options,
+            dataset_name=dataset_name,
+            dataset_version=dataset_version
         )
-        stats | 'Write stats' >> tfdv.WriteStatisticsToTFRecord(
-            os.path.join(output_path, 'stats.tfrecord'))
-
-        # Write out examples
-        (training_examples
-         | 'Write training' >> transforms.WriteTFExample(
-             os.path.join(output_path, 'training'))
+    elif output_type == 'tif' or output_type == 'tiff':
+        from geebeam import tiff_writer
+        tiff_writer.run_tiff_export(
+            input_records=input_records,
+            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
         )
-        if config['validation_ratio'] > 0:
-            validation_examples = (
-                validation_data
-                | 'Val to tf.Example' >> beam.Map(transforms.dict_to_example)
-            )
-
-            (validation_examples
-            | 'Write validation' >> transforms.WriteTFExample(
-                os.path.join(output_path, 'validation'))
-            )
-
-
-    # NOTE: testing a few different formats.
-    # Need to simplify once we decide on one
-    # Infer schema
-    stats = tfdv.load_statistics(
-        os.path.join(output_path, 'stats.tfrecord')
-    )
-
-    schema = tfdv.infer_schema(stats)
-
-    tfdv.write_schema_text(
-        schema,
-        os.path.join(output_path, 'schema.pbtxt')
-    )
-
-    # Also write as pbtxt, easier to read
-    tfdv.write_stats_text(
-        stats,
-        os.path.join(output_path, 'stats.pbtxt')
-    )
-
-    # Also write stats and schema as jsons
-    out_schema_json = os.path.join(output_path, 'schema.json')
-    out_stats_json = os.path.join(output_path, 'stats.json')
-    if output_path.startswith('gs://'):
-        transforms.write_json_to_gcs(MessageToJson(schema), out_schema_json)
-        transforms.write_json_to_gcs(MessageToJson(stats), out_stats_json)
     else:
-        transforms.write_json_to_local(MessageToJson(schema), out_schema_json)
-        transforms.write_json_to_local(MessageToJson(stats), out_stats_json)
-
-
+        raise ValueError(f"output_type {output_type} not implemented. Options are ['tfds', 'tfrecord', 'tiff']")

geebeam-0.3.2/src/geebeam/tf_utils.py

@@ -0,0 +1,73 @@
+"""Helpers for tensorflow"""
+
+import tensorflow as tf
+import apache_beam as beam
+
+from geebeam import transforms
+
+def _bytes_feature(value):
+    """Returns a bytes_list from a string / byte."""
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
+
+def _float_feature(value):
+    """Returns a float_list from a float / double."""
+    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
+
+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):
+    """"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=
+                tf.train.FloatList(
+                    value = transforms.convert_to_iterable(element['metadata'][md_key])
+                )
+            )
+
+    # Build image feature with named bands
+    array_dict = {}
+    for im_feat in element['array'].keys():#.dtype.names:
+        array_dict['im_'+im_feat] = tf.train.Feature(
+            float_list = tf.train.FloatList(
+                value = element['array'][im_feat].flatten()))
+
+    # Combine
+    feature = {**md_dict, **array_dict}
+
+    # Build example and serialize
+    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'):
+        self.output_dir = output_dir
+        self.file_name_suffix = file_name_suffix
+
+    def expand(self, pcoll):
+        return (
+            pcoll
+            | 'Write to TFRecord' >> beam.io.WriteToTFRecord(
+                self.output_dir,
+                file_name_suffix=self.file_name_suffix
+            )
+        )