Amazon SageMaker 地理空间功能入门

import boto3
import sagemaker
import sagemaker_geospatial_map

import time
import datetime
import os
from glob import glob
import rasterio
from rasterio.plot import show
import matplotlib.colors
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import tifffile

sagemaker_session = sagemaker.Session()
export_bucket = sagemaker_session.default_bucket() # Alternatively you can use your custom bucket here. 

session = boto3.Session()
execution_role = sagemaker.get_execution_role()
geospatial_client = session.client(service_name="sagemaker-geospatial")

eoj_input_config = {
    "RasterDataCollectionQuery": {
        "RasterDataCollectionArn": "arn:aws:sagemaker-geospatial:us-west-2:378778860802:raster-data-collection/public/nmqj48dcu3g7ayw8",
        "AreaOfInterest": {
            "AreaOfInterestGeometry": {
                "PolygonGeometry": {
                    "Coordinates": [
                        [
                            [-121.32559295351282, 40.386534879495315],
                            [-121.32559295351282, 40.09770246706907],
                            [-120.86738632168885, 40.09770246706907],
                            [-120.86738632168885, 40.386534879495315],
                            [-121.32559295351282, 40.386534879495315]
                        ]
                    ]
                }
            }
        },
        "TimeRangeFilter": {
            "StartTime": "2021-06-01T00:00:00Z",
            "EndTime": "2021-09-30T23:59:59Z",
        },
        "PropertyFilters": {
            "Properties": [{"Property": {"EoCloudCover": {"LowerBound": 0, "UpperBound": 0.1}}}],
            "LogicalOperator": "AND",
        },
    }
}

eoj_config = {"LandCoverSegmentationConfig": {}}

response = geospatial_client.start_earth_observation_job(
    Name="dixie-wildfire-landcover-2021",
    InputConfig=eoj_input_config,
    JobConfig=eoj_config,
    ExecutionRoleArn=execution_role,
)
eoj_arn = response["Arn"]
eoj_arn 

search_params = eoj_input_config.copy()
search_params["Arn"] = "arn:aws:sagemaker-geospatial:us-west-2:378778860802:raster-data-collection/public/nmqj48dcu3g7ayw8"
search_params["RasterDataCollectionQuery"].pop("RasterDataCollectionArn", None)
search_params["RasterDataCollectionQuery"]["BandFilter"] = ["visual"]

cog_urls = []
search_result = geospatial_client.search_raster_data_collection(**search_params)
for item in search_result["Items"]:
    asset_url = item["Assets"]["visual"]["Href"]
    cog_urls.append(asset_url)

cog_urls

cog_urls.sort(key=lambda x: x.split("TFK_")[1])

src_pre = rasterio.open(cog_urls[0])
src_post = rasterio.open(cog_urls[-1])

fig, (ax_before, ax_after) = plt.subplots(1, 2, figsize=(14,7))
subplot = show(src_pre, ax=ax_before)
subplot.axis('off')
subplot.set_title("Pre-wildfire ({})".format(cog_urls[0].split("TFK_")[1]))
subplot = show(src_post, ax=ax_after)
subplot.axis('off')
subplot.set_title("Post-wildfire ({})".format(cog_urls[-1].split("TFK_")[1]))
plt.show()

# check status of created Earth Observation Job and wait until it is completed
eoj_completed = False
while not eoj_completed:
    response = geospatial_client.get_earth_observation_job(Arn=eoj_arn)
    print("Earth Observation Job status: {} (Last update: {})".format(response['Status'], datetime.datetime.now()), end='\r')
    eoj_completed = True if response['Status'] == 'COMPLETED' else False
    if not eoj_completed:
        time.sleep(30)

bucket_prefix = "eoj_dixie_wildfire_landcover"
response = geospatial_client.export_earth_observation_job(
    Arn=eoj_arn,
    ExecutionRoleArn=execution_role,
    OutputConfig={
        "S3Data": {"S3Uri": f"s3://{export_bucket}/{bucket_prefix}/"}
    },
)

while not response['ExportStatus'] == 'SUCCEEDED':
    response = geospatial_client.get_earth_observation_job(Arn=eoj_arn)
    print("Export of Earth Observation Job status: {} (Last update: {})".format(response['ExportStatus'], datetime.datetime.now()), end='\r')
    if not response['ExportStatus'] == 'SUCCEEDED':
        time.sleep(30)

s3_bucket = session.resource("s3").Bucket(export_bucket)

mask_dir = "./dixie-wildfire-landcover/masks"
os.makedirs(mask_dir, exist_ok=True)
for s3_object in s3_bucket.objects.filter(Prefix=bucket_prefix).all():
    path, filename = os.path.split(s3_object.key)
    if "output" in path:
        mask_local_path = mask_dir + "/" + filename
        s3_bucket.download_file(s3_object.key, mask_local_path)
        print("Downloaded mask: " + mask_local_path)

mask_files = glob(os.path.join(mask_dir, "*.tif"))
mask_files.sort(key=lambda x: x.split("TFK_")[1])

landcover_simple_colors = {"not vegetated": "khaki","vegetated": "olivedrab", "water": "lightsteelblue"}

def extract_masks(date_str):
    mask_file = list(filter(lambda x: date_str in x, mask_files))[0]
    mask = tifffile.imread(mask_file)
    focus_area_mask = mask[400:1100, 600:1350]
    
    vegetation_mask = np.isin(focus_area_mask, [4]).astype(np.uint8) # vegetation has a class index of 4
    water_mask = np.isin(focus_area_mask, [6]).astype(np.uint8) # water has a class index of 6
    water_mask[water_mask > 0] = 2
    additive_mask = np.add(vegetation_mask, water_mask).astype(np.uint8)
    
    return (focus_area_mask, vegetation_mask, additive_mask)

masks_20210603 = extract_masks("20210603")
masks_20210926 = extract_masks("20210926")

fig = plt.figure(figsize=(14,7))

fig.add_subplot(1, 2, 1)
plt.imshow(masks_20210603[2], cmap=matplotlib.colors.ListedColormap(list(landcover_simple_colors.values()), N=None))
plt.title("Pre-wildfire")
plt.axis('off')
ax = fig.add_subplot(1, 2, 2)
hs = plt.imshow(masks_20210926[2], cmap=matplotlib.colors.ListedColormap(list(landcover_simple_colors.values()), N=None))
plt.title("Post-wildfire")
plt.axis('off')
patches = [ mpatches.Patch(color=i[1], label=i[0]) for i in landcover_simple_colors.items()]
plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0. )
plt.show()

vegetation_loss = round((1 - (masks_20210926[1].sum() / masks_20210603[1].sum())) * 100, 2)
diff_mask = np.add(masks_20210603[1], masks_20210926[1])
plt.figure(figsize=(6, 6))
plt.title("Loss in vegetation ({}%)".format(vegetation_loss))
plt.imshow(diff_mask, cmap=matplotlib.colors.ListedColormap(["black","crimson", "silver"], N=None))
plt.axis('off')
patches = [mpatches.Patch(color="crimson", label="vegetation lost")]
plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0. )
plt.show()