croplands/io.py

import xarray as xr
from pathlib import Path
import numpy as np
from affine import Affine
from rasterio import CRS


def from_spatial_to_profile(zarr):
  profile = {}
  profile["driver"] = "GTiff"
  profile["dtype"] =  zarr.dtypes["band_data"].name
  profile["width"] = zarr.dims["x"]
  profile["height"] = zarr.dims["y"]
  profile["count"] = zarr.dims["band"]
  profile["crs"] = CRS.from_epsg(zarr["spatial_ref"].attrs["crs_wkt"].split(",")[-1].split("\"")[1])
  vals = zarr["spatial_ref"].attrs["GeoTransform"].split(" ")
  profile["transform"] = Affine(float(vals[1]),float(vals[2]),float(vals[0]),float(vals[4]),float(vals[5]),float(vals[3]))
  return profile


def read_zarr(
    path: Path
) -> np.ndarray:
    """Read a zarr file using xarray.

    Args:
        path (Path): Path to the zarr file.
    Returns:
        np.ndarray: zarr data as numpy array.
    """
    data = xr.open_zarr(path)["band_data"]
    return data.values

def read_zarr_profile(path:Path):
    data = xr.open_zarr(path)
    profile = from_spatial_to_profile(data)
    return profile