Upload 3 files

chmloader.py

@@ -0,0 +1,287 @@
+import os
+import geopandas as gpd
+from osgeo import gdal
+import rasterio
+from rasterio.mask import mask
+import numpy as np
+from rasterio.warp import calculate_default_transform, reproject, Resampling
+from pyproj import CRS
+
+def build_chm_srcs(target, tiles_zip='tiles/tiles.zip', local_chm_dir='chm_tiles'):
+    """
+    Get the S3 paths to the CHM tiles that cover a given target area. Download missing tiles from S3 if necessary.
+
+    Args:
+        target: GeoDataFrame or GeoSeries representing the target area.
+        tiles_zip: Path to the zip file containing the tiles.
+        local_chm_dir: Directory where CHM tiles are stored locally.
+
+    Returns:
+        list: A list of paths to the CHM tiles that cover the target area.
+    """
+    # Load the tiles from the zip file
+    gtiles = gpd.read_file(f'/vsizip/{tiles_zip}')
+
+    # Transform the target area to match the tiles' CRS
+    t_ext = target.to_crs(gtiles.crs).total_bounds
+
+    # Filter the tiles that intersect with the target area
+    filtered_tiles = gtiles.cx[t_ext[0]:t_ext[2], t_ext[1]:t_ext[3]]
+
+    # Construct paths to the CHM files using GDAL's /vsis3/ to access S3 without credentials
+    s3_paths = [
+        f'/vsis3/dataforgood-fb-data/forests/v1/alsgedi_global_v6_float/chm/{tile}.tif' for tile in filtered_tiles['tile']
+    ]
+
+    return s3_paths
+
+def warp_util(srcs, target, res, filename, gdalwarp_options, gdal_config_options):
+    """
+    A gdalwarp utility function.
+
+    Args:
+        srcs (list): The source files to warp.
+        target: GeoDataFrame or GeoSeries representing the target area to warp to.
+        res (float or tuple): The resolution of the target area.
+        filename (str): The filename to save the warped data to.
+        gdalwarp_options (list): Options to pass to gdalwarp.
+        gdal_config_options (dict): GDAL configuration options to set before running gdalwarp.
+
+    Returns:
+        str: The path to the warped file.
+    """
+    # Set GDAL configuration options
+    for key, value in gdal_config_options.items():
+        gdal.SetConfigOption(key, value)
+
+    # Calculate resolution if not provided
+    if res is None:
+        bbox = target.total_bounds
+        xmin, ymin, xmax, ymax = bbox
+        res_x = 1.0
+        res_y = 1.0
+        dim_x = int((xmax - xmin) / res_x)
+        dim_y = int((ymax - ymin) / res_y)
+    else:
+        if isinstance(res, (int, float)):
+            res_x = res
+            res_y = res
+        elif isinstance(res, (tuple, list)) and len(res) == 2:
+            res_x, res_y = res
+        else:
+            raise ValueError("Resolution must be a single value or a tuple of two values.")
+        bbox = target.total_bounds
+        xmin, ymin, xmax, ymax = bbox
+        dim_x = int((xmax - xmin) / res_x)
+        dim_y = int((ymax - ymin) / res_y)
+
+    te_srs = target.crs.to_string()
+
+    # Add translation extent, target SRS, and output dimensions to gdalwarp options
+    gdalwarp_options += [
+        "-te", str(xmin), str(ymin), str(xmax), str(ymax),
+        "-t_srs", te_srs,
+        "-ts", str(dim_x), str(dim_y)
+    ]
+
+    # Perform the warp operation
+    gdal.Warp(
+        filename, srcs, options=gdalwarp_options
+    )
+
+    # Reset GDAL configuration options to their defaults after the operation
+    for key in gdal_config_options.keys():
+        gdal.SetConfigOption(key, None)
+
+    return filename
+
+def download_chm(target, chm="tolan", res=None, filename="output_chm.tif", gdalwarp_options=None, gdal_config_options=None):
+    """
+    Download Canopy Height Model (CHM) data by mosaicking relevant CHM tiles.
+
+    Args:
+        target: GeoDataFrame or GeoSeries representing the target area.
+        chm (str): The CHM dataset to download (currently only 'tolan').
+        res (float): The resolution of the CHM data to download in meters.
+        filename (str): The filename to save the downloaded CHM data to.
+        gdalwarp_options (list): Options to pass to gdalwarp.
+        gdal_config_options (dict): GDAL configuration options to set before running gdalwarp.
+
+    Returns:
+        str: The path to the downloaded CHM data.
+    """
+    print("downloading CHM")
+    if gdalwarp_options is None:
+        gdalwarp_options = ["-r", "bilinear", "-overwrite"]
+    if gdal_config_options is None:
+        gdal_config_options = {
+            "GDAL_CACHEMAX": "512",
+            "GDAL_HTTP_MULTIPLEX": "YES",
+            "VSI_CACHE": "TRUE",
+            "AWS_NO_SIGN_REQUEST": "YES"
+        }
+
+    # Generate the paths to the source CHM tiles
+    srcs = build_chm_srcs(target)
+
+    # Warp and mosaic the CHM tiles into a single output file
+    chm_raster_path = warp_util(srcs, target, res, filename, gdalwarp_options, gdal_config_options)
+
+    return chm_raster_path
+
+
+# Ensure the NoData Handling is Correct
+# def clip_and_remove_zeros(tif_path, boundary, output_tif_path=None, nodata_value=None, save=False):
+#     with rasterio.open(tif_path) as src:
+#         out_image, out_transform = mask(src, boundary.geometry, crop=True)
+#         profile = src.profile.copy()
+
+#         # Determine the NoData value
+#         if nodata_value is None:
+#             if src.nodata is not None:
+#                 nodata_value = src.nodata
+#             else:
+#                 nodata_value = 0  # Set default NoData value
+
+#         # Mask out NoData values
+#         out_image = out_image.squeeze()
+#         out_image = np.ma.masked_equal(out_image, nodata_value)
+
+#         # Update profile for saving, if needed
+#         profile.update({
+#             'dtype': src.dtypes[0],
+#             'count': 1,
+#             'nodata': nodata_value,
+#             'compress': 'deflate'
+#         })
+
+#         if save and output_tif_path:
+#             with rasterio.open(output_tif_path, 'w', **profile) as dst:
+#                 dst.write(out_image, 1)
+
+#     return out_image, out_transform, profile
+
+
+def reproject_to_utm(tif_path):
+    """Reproject the raster file to UTM and save it with the same name."""
+    with rasterio.open(tif_path) as src:
+        # Get the CRS of the file
+        crs = src.crs
+        print(f"Original CRS of the TIFF file: {crs}")
+
+        if crs.is_geographic:
+            # Calculate the center point of the raster
+            center_lon, center_lat = (src.bounds.left + src.bounds.right) / 2, (src.bounds.bottom + src.bounds.top) / 2
+            print(f"Center of the raster: Longitude={center_lon}, Latitude={center_lat}")
+
+            # Determine UTM zone
+            utm_zone = int((center_lon + 180) // 6) + 1
+            hemisphere = 'north' if center_lat >= 0 else 'south'
+
+            # Define the UTM CRS
+            if hemisphere == 'north':
+                epsg_code = 32600 + utm_zone  # UTM zone for Northern Hemisphere
+            else:
+                epsg_code = 32700 + utm_zone  # UTM zone for Southern Hemisphere
+
+            utm_crs = CRS.from_epsg(epsg_code)
+            print(f"Determined UTM CRS: EPSG:{epsg_code}")
+
+        elif crs.is_projected:
+            # If CRS is already projected, attempt to retrieve its EPSG code
+            try:
+                epsg_code = crs.to_epsg()  # Extract EPSG code
+                if epsg_code is not None:
+                    print(f"The CRS is already projected. EPSG code: {epsg_code}")
+                    utm_crs = CRS.from_epsg(epsg_code)
+                else:
+                    # Fallback to using the WKT representation if EPSG is None
+                    utm_crs = crs
+                    print(f"The CRS is projected but has no EPSG code. Using WKT: {crs.to_wkt()}")
+            except Exception as e:
+                print(f"Unable to determine EPSG code or WKT: {e}")
+                utm_crs = crs
+        else:
+            print("The CRS type could not be determined or is not geographic.")
+            return tif_path  # No reprojection needed, return the original path
+
+        # Prepare output path (same as input to overwrite)
+        output_path = tif_path
+
+        # Calculate the transform and dimensions for the reprojected image
+        transform, width, height = calculate_default_transform(
+            src.crs, utm_crs, src.width, src.height, *src.bounds
+        )
+
+        # Define the output profile for the new UTM raster
+        profile = src.profile.copy()
+        profile.update({
+            'crs': utm_crs,
+            'transform': transform,
+            'width': width,
+            'height': height
+        })
+
+        # Reproject the raster to UTM
+        with rasterio.open(output_path, 'w', **profile) as dst:
+            for i in range(1, src.count + 1):
+                reproject(
+                    source=rasterio.band(src, i),
+                    destination=rasterio.band(dst, i),
+                    src_transform=src.transform,
+                    src_crs=src.crs,
+                    dst_transform=transform,
+                    dst_crs=utm_crs,
+                    resampling=Resampling.nearest
+                )
+
+    print(f"Reprojected file saved as: {output_path}")
+    return output_path  # Return the path of the reprojected file
+
+
+def clip_and_remove_zeros(tif_path, boundary, output_tif_path=None, nodata_value=None, save=False):
+    # Reproject the raster to UTM if necessary
+    reprojected_path = reproject_to_utm(tif_path)
+
+    # Open the reprojected raster file
+    with rasterio.open(reprojected_path) as src:
+        # Check if the CRS is UTM
+        if src.crs.is_projected and src.crs.to_epsg() and (32601 <= src.crs.to_epsg() <= 32660 or 32701 <= src.crs.to_epsg() <= 32760):
+            print(f"The raster is in UTM CRS: EPSG:{src.crs.to_epsg()}")
+        else:
+            print("Warning: The raster is not in a UTM CRS or its CRS could not be determined to be UTM.")
+
+        # Ensure the boundary is in the same CRS as the raster
+        if boundary.crs != src.crs:
+            print(f"Reprojecting boundary to match raster CRS: {src.crs}")
+            boundary = boundary.to_crs(src.crs)
+        
+        # Clip the raster using the boundary geometry
+        out_image, out_transform = mask(src, boundary.geometry, crop=True)
+        profile = src.profile.copy()
+
+        # Determine the NoData value
+        if nodata_value is None:
+            if src.nodata is not None:
+                nodata_value = src.nodata
+            else:
+                nodata_value = 0  # Set default NoData value
+
+        # Mask out NoData values
+        out_image = out_image.squeeze()
+        out_image = np.ma.masked_equal(out_image, nodata_value)
+
+        # Update profile for saving, if needed
+        profile.update({
+            'dtype': src.dtypes[0],
+            'count': 1,
+            'nodata': nodata_value,
+            'compress': 'deflate'
+        })
+
+        if save and output_tif_path:
+            with rasterio.open(output_tif_path, 'w', **profile) as dst:
+                dst.write(out_image, 1)
+
+    return out_image, out_transform, profile
+