From 5f410fcfd58b825b8f26241d972fdf451d6d51e6 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 24 Apr 2024 23:06:43 -0700 Subject: [PATCH 01/11] WIP run script for patch level world cover embeddings generation --- scripts/worldcover/run.py | 237 ++++++++++++++++++++++++++++++++------ 1 file changed, 201 insertions(+), 36 deletions(-) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index d1458e1b..ba40e9b9 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -1,19 +1,24 @@ #!/usr/bin/env python3 -# import sys -# sys.path.append("/home/tam/Documents/repos/model") +import sys +sys.path.append("/home/ubuntu/worldcover/model") import os import tempfile from math import floor +from pathlib import Path +import requests import boto3 import einops import geopandas as gpd +import pandas as pd import numpy import pyarrow as pa import rasterio +import shapely import torch +import xarray as xr from rasterio.windows import Window from shapely import box from torchvision.transforms import v2 @@ -24,6 +29,7 @@ YEAR = int(os.environ.get("YEAR", 2020)) DATE = f"{YEAR}-06-01" TILE_SIZE = 12000 +PATCH_SIZE = 32 CHIP_SIZE = 512 E_W_INDEX_START = 67 E_W_INDEX_END = 125 @@ -36,9 +42,12 @@ RASTER_X_SIZE = (E_W_INDEX_END - E_W_INDEX_START) * TILE_SIZE RASTER_Y_SIZE = (N_S_INDEX_END - N_S_INDEX_START) * TILE_SIZE NODATA = 0 -CKPT_PATH = "s3://clay-model-ckpt/v0/mae_epoch-24_val-loss-0.46.ckpt" +CKPT_PATH = ( + "https://huggingface.co/made-with-clay/Clay/resolve/main/" + "Clay_v0.1_epoch-24_val-loss-0.46.ckpt" +) # CKPT_PATH = "https://huggingface.co/made-with-clay/Clay/resolve/main/Clay_v0.1_epoch-24_val-loss-0.46.ckpt" -VERSION = "002" +VERSION = "003" BUCKET = "clay-worldcover-embeddings" URL = "https://esa-worldcover-s2.s3.amazonaws.com/rgbnir/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}W{xidx}_S2RGBNIR.tif" WC_VERSION_LOOKUP = { @@ -134,6 +143,62 @@ def tiles_and_windows(input: Window): return result +def download_image(url): + # Download the image from the URL + response = requests.get(url) + # Check if the request was successful + if response.status_code == 200: + return response.content # Return the image content + else: + raise Exception("Failed to download the image") + +def patches_and_windows_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): + # Download the image from the URL + image_data = download_image(url) + + # Open the image using rasterio from memory + with rasterio.io.MemoryFile(image_data) as memfile: + with memfile.open() as src: + # Read the image data and metadata + img_data = src.read() + img_meta = src.profile + img_crs = src.crs + + # Convert raster data and metadata into an xarray DataArray + img_da = xr.DataArray(img_data, dims=("band", "y", "x"), attrs=img_meta) + + # Tile the data + ds_chunked = img_da.chunk({"y": chunk_size[0], "x": chunk_size[1]}) + + # Get the geospatial information from the original dataset + transform = img_meta["transform"] + + # Iterate over the chunks and compute the geospatial bounds for each chunk + chunk_bounds = {} + + for x in range(ds_chunked.sizes["x"] // chunk_size[1]): + for y in range(ds_chunked.sizes["y"] // chunk_size[0]): + # Compute chunk coordinates + x_start = x * chunk_size[1] + y_start = y * chunk_size[0] + x_end = min(x_start + chunk_size[1], ds_chunked.sizes["x"]) + y_end = min(y_start + chunk_size[0], ds_chunked.sizes["y"]) + + # Compute chunk geospatial bounds + lon_start, lat_start = transform * (x_start, y_start) + lon_end, lat_end = transform * (x_end, y_end) + + # Store chunk bounds + chunk_bounds[(x, y)] = { + "lon_start": lon_start, + "lat_start": lat_start, + "lon_end": lon_end, + "lat_end": lat_end, + } + + return chunk_bounds, img_crs + + def make_batch(result): pixels = [] for url, win in result: @@ -168,8 +233,36 @@ def make_batch(result): "timestep": torch.as_tensor(data=[ds.normalize_timestamp(f"{YEAR}-06-01")]).to( rgb_model.device ), + "date": f"{YEAR}-06-01" + , } +def get_pixels(result): + pixels = [] + for url, win in result: + with rasterio.open(url) as src: + data = src.read(window=win) + if NODATA in data: + return + pixels.append(data) + transform = src.window_transform(win) + + if len(pixels) == 1: + pixels = pixels[0] + elif len(pixels) == 2: # noqa: PLR2004 + if pixels[0].shape[2] == CHIP_SIZE: + pixels = einops.pack(pixels, "b * w")[0] + else: + pixels = einops.pack(pixels, "b h *")[0] + else: + px1 = einops.pack(pixels[:2], "b w *")[0] + px2 = einops.pack(pixels[2:], "b w *")[0] + pixels = einops.pack((px1, px2), "b * w")[0] + + assert pixels.shape == (4, CHIP_SIZE, CHIP_SIZE) + + return pixels + index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 2)) @@ -177,19 +270,36 @@ def make_batch(result): tfm = v2.Compose([v2.Normalize(mean=MEAN, std=STD)]) ds = ClayDataset(chips_path=[], transform=tfm) +CKPT_PATH = ( + "https://huggingface.co/made-with-clay/Clay/resolve/main/" + "Clay_v0.1_epoch-24_val-loss-0.46.ckpt" +) + +# Load model rgb_model = CLAYModule.load_from_checkpoint( CKPT_PATH, mask_ratio=0.0, - band_groups={"rgb": (0, 1, 2), "nir": (3,)}, + band_groups={"rgb": (2, 1, 0), "nir": (3,)}, + bands=4, strict=False, # ignore the extra parameters in the checkpoint + embeddings_level="group", ) +# Set the model to evaluation mode +rgb_model.eval() + + +outdir_embeddings = Path("data/embeddings") +outdir_embeddings.mkdir(exist_ok=True, parents=True) xoff = index * CHIP_SIZE yoff = 0 embeddings = [] all_bounds = [] +results = [] while yoff < RASTER_Y_SIZE: result = tiles_and_windows(Window(xoff, yoff, CHIP_SIZE, CHIP_SIZE)) + if result is not None: + results.append(result) if result is None: yoff += CHIP_SIZE @@ -218,34 +328,89 @@ def make_batch(result): ) yoff += CHIP_SIZE - -embeddings = numpy.vstack(embeddings) - -embeddings_mean = embeddings[:, :-2, :].mean(axis=1) - -print(f"Average embeddings have shape {embeddings_mean.shape}") - -gdf = gpd.GeoDataFrame( - data={ - "embeddings": pa.FixedShapeTensorArray.from_numpy_ndarray( - numpy.ascontiguousarray(embeddings_mean) - ), - }, - geometry=[box(*dat) for dat in all_bounds], # This assumes same order - crs="EPSG:4326", -) - -with tempfile.TemporaryDirectory() as tmp: - # tmp = "/home/tam/Desktop/wcctmp" - - outpath = f"{tmp}/worldcover_embeddings_{YEAR}_{index}_v{VERSION}.gpq" - print(f"Uploading embeddings to {outpath}") - - gdf.to_parquet(path=outpath, compression="ZSTD", schema_version="1.0.0") - - s3_client = boto3.client("s3") - s3_client.upload_file( - outpath, - BUCKET, - f"v{VERSION}/{YEAR}/{os.path.basename(outpath)}", - ) + + + + print(len(embeddings), len(results)) + #embeddings = numpy.vstack(embeddings) + embeddings_ = embeddings[0] + print("Embeddings shape: ", embeddings_.shape) + + embeddings_ = embeddings_[:, :-2, :] + + print(f"Embeddings have shape {embeddings_.shape}") #.mean(axis=1) + + # remove date and lat/lon and reshape to disaggregated patches + embeddings_patch = embeddings_.reshape([2, 16, 16, 768]) + + # average over the band groups + embeddings_mean = embeddings_patch.mean(axis=0) + + print(f"Average patch embeddings have shape {embeddings_mean.shape}") + + + if result is not None: + print("result: ", result[0][0]) + pix = get_pixels(result) + chunk_bounds, epsg = patches_and_windows_from_url(result[0][0]) + #print("chunk_bounds: ", chunk_bounds) + print("chunk bounds length:", len(chunk_bounds)) + + # Iterate through each patch + for i in range(embeddings_mean.shape[0]): + for j in range(embeddings_mean.shape[1]): + embeddings_output_patch = embeddings_mean[i, j] + + item_ = [ + element for element in list(chunk_bounds.items()) if element[0] == (i, j) + ] + box_ = [ + item_[0][1]["lon_start"], + item_[0][1]["lat_start"], + item_[0][1]["lon_end"], + item_[0][1]["lat_end"], + ] + #source_url = batch["source_url"] + date = batch["date"] + date_as_timestamp = pd.to_datetime(date, format="%Y-%m-%d") + + # Convert the Pandas Timestamp to the desired data type + #date_as_date32 = date_as_timestamp.astype('datetime64[D]') + + #print(batch["date"]) + data = { + "date": date_as_timestamp, + "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], + } + + # Define the bounding box as a Polygon (xmin, ymin, xmax, ymax) + # The box_ list is encoded as + # [bottom left x, bottom left y, top right x, top right y] + box_emb = shapely.geometry.box(box_[0], box_[1], box_[2], box_[3]) + + print(str(epsg)[-4:]) + + # Create the GeoDataFrame + gdf = gpd.GeoDataFrame(data, geometry=[box_emb], crs=f"EPSG:{str(epsg)[-4:]}") + + # Reproject to WGS84 (lon/lat coordinates) + gdf = gdf.to_crs(epsg=4326) + + + with tempfile.TemporaryDirectory() as tmp: + # tmp = "/home/tam/Desktop/wcctmp" + + outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_v{VERSION}.gpq" + print(f"Uploading embeddings to {outpath}") + #print(gdf) + + gdf.to_parquet(path=outpath, compression="ZSTD", schema_version="1.0.0") + + s3_client = boto3.client("s3") + s3_client.upload_file( + outpath, + BUCKET, + f"v{VERSION}/{YEAR}/{os.path.basename(outpath)}", + ) + + \ No newline at end of file From 081e8de35b428e9fd8a5eb7aea21b6fdce9f0b25 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 24 Apr 2024 23:09:57 -0700 Subject: [PATCH 02/11] remove dup lines --- scripts/worldcover/run.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index ba40e9b9..fbbc4708 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -270,11 +270,6 @@ def get_pixels(result): tfm = v2.Compose([v2.Normalize(mean=MEAN, std=STD)]) ds = ClayDataset(chips_path=[], transform=tfm) -CKPT_PATH = ( - "https://huggingface.co/made-with-clay/Clay/resolve/main/" - "Clay_v0.1_epoch-24_val-loss-0.46.ckpt" -) - # Load model rgb_model = CLAYModule.load_from_checkpoint( CKPT_PATH, From d9e8eb2977fba40e5b46e6e1c928062b7792ee65 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 24 Apr 2024 23:18:07 -0700 Subject: [PATCH 03/11] clean up paths --- scripts/worldcover/embeddings_db.py | 14 +++++++------- scripts/worldcover/run.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/scripts/worldcover/embeddings_db.py b/scripts/worldcover/embeddings_db.py index 1f106448..bb53ae02 100644 --- a/scripts/worldcover/embeddings_db.py +++ b/scripts/worldcover/embeddings_db.py @@ -6,12 +6,12 @@ from skimage import io # Set working directory -wd = "/home/usr/Desktop/" +wd = "./" # To download the existing embeddings run aws s3 sync # aws s3 sync s3://clay-worldcover-embeddings /my/dir/clay-worldcover-embeddings -vector_dir = Path(wd + "clay-worldcover-embeddings/v002/2021/") +vector_dir = Path(wd + "clay-worldcover-embeddings/2020/") # Create new DB structure or open existing db = lancedb.connect(wd + "worldcoverembeddings_db") @@ -24,17 +24,17 @@ for _, row in tile_df.iterrows(): data.append( - {"vector": row["embeddings"], "year": 2021, "bbox": row.geometry.bounds} + {"vector": row["embeddings"], "year": 2020, "bbox": row.geometry.bounds} ) # Show table names db.table_names() # Drop existing table if exists -db.drop_table("worldcover-2021-v001") +#db.drop_table("worldcover-2020-v001") # Create embeddings table and insert the vector data -tbl = db.create_table("worldcover-2021-v001", data=data, mode="overwrite") +tbl = db.create_table("worldcover-2020-v001", data=data, mode="overwrite") # Visualize some image chips @@ -53,6 +53,6 @@ def plot(df, cols=10): # Select a vector by index, and search 10 similar pairs, and plot -v = tbl.to_pandas()["vector"].values[10540] +v = tbl.to_pandas()["vector"].values[5] result = tbl.search(query=v).limit(5).to_pandas() -plot(result, 5) +plot(result, 5) \ No newline at end of file diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index fbbc4708..575ddfd5 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import sys -sys.path.append("/home/ubuntu/worldcover/model") +sys.path.append("../../") import os import tempfile From 7bb7909ab88bad18844ad43b8c41043b27085f64 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Fri, 26 Apr 2024 09:36:23 -0700 Subject: [PATCH 04/11] restructuring --- scripts/worldcover/run.py | 42 +++++++++++++++++---------------------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index 575ddfd5..95ad71b5 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -46,8 +46,7 @@ "https://huggingface.co/made-with-clay/Clay/resolve/main/" "Clay_v0.1_epoch-24_val-loss-0.46.ckpt" ) -# CKPT_PATH = "https://huggingface.co/made-with-clay/Clay/resolve/main/Clay_v0.1_epoch-24_val-loss-0.46.ckpt" -VERSION = "003" +VERSION = "005" BUCKET = "clay-worldcover-embeddings" URL = "https://esa-worldcover-s2.s3.amazonaws.com/rgbnir/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}W{xidx}_S2RGBNIR.tif" WC_VERSION_LOOKUP = { @@ -142,9 +141,8 @@ def tiles_and_windows(input: Window): return result - def download_image(url): - # Download the image from the URL + # Download an image from a URL response = requests.get(url) # Check if the request was successful if response.status_code == 200: @@ -152,8 +150,8 @@ def download_image(url): else: raise Exception("Failed to download the image") -def patches_and_windows_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): - # Download the image from the URL +def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): + # Download an image from a URL image_data = download_image(url) # Open the image using rasterio from memory @@ -198,7 +196,6 @@ def patches_and_windows_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): return chunk_bounds, img_crs - def make_batch(result): pixels = [] for url, win in result: @@ -282,7 +279,6 @@ def get_pixels(result): # Set the model to evaluation mode rgb_model.eval() - outdir_embeddings = Path("data/embeddings") outdir_embeddings.mkdir(exist_ok=True, parents=True) @@ -327,15 +323,16 @@ def get_pixels(result): print(len(embeddings), len(results)) - #embeddings = numpy.vstack(embeddings) - embeddings_ = embeddings[0] + embeddings_ = numpy.vstack(embeddings) + #embeddings_ = embeddings[0] print("Embeddings shape: ", embeddings_.shape) + + # remove date and lat/lon + embeddings_ = embeddings_[:, :-2, :].mean(axis=0) - embeddings_ = embeddings_[:, :-2, :] - - print(f"Embeddings have shape {embeddings_.shape}") #.mean(axis=1) + print(f"Embeddings have shape {embeddings_.shape}") - # remove date and lat/lon and reshape to disaggregated patches + # reshape to disaggregated patches embeddings_patch = embeddings_.reshape([2, 16, 16, 768]) # average over the band groups @@ -347,7 +344,7 @@ def get_pixels(result): if result is not None: print("result: ", result[0][0]) pix = get_pixels(result) - chunk_bounds, epsg = patches_and_windows_from_url(result[0][0]) + chunk_bounds, epsg = patch_bounds_from_url(result[0][0]) #print("chunk_bounds: ", chunk_bounds) print("chunk bounds length:", len(chunk_bounds)) @@ -365,16 +362,14 @@ def get_pixels(result): item_[0][1]["lon_end"], item_[0][1]["lat_end"], ] - #source_url = batch["source_url"] - date = batch["date"] - date_as_timestamp = pd.to_datetime(date, format="%Y-%m-%d") - - # Convert the Pandas Timestamp to the desired data type - #date_as_date32 = date_as_timestamp.astype('datetime64[D]') - #print(batch["date"]) data = { - "date": date_as_timestamp, + #"source_url": batch["source_url"][0], + #"date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( + # dtype="date32[day][pyarrow]" + #), + #"date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), + "date": pd.to_datetime(batch["date"], format="%Y-%m-%d"), "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], } @@ -390,7 +385,6 @@ def get_pixels(result): # Reproject to WGS84 (lon/lat coordinates) gdf = gdf.to_crs(epsg=4326) - with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" From a7dd3c8bc16095494e99bb52e9eb5919a6fe30d9 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 1 May 2024 12:12:55 -0700 Subject: [PATCH 05/11] add SWIR composite --- scripts/worldcover/run_msi.py | 419 ++++++++++++++++++++++++++++++++++ 1 file changed, 419 insertions(+) create mode 100644 scripts/worldcover/run_msi.py diff --git a/scripts/worldcover/run_msi.py b/scripts/worldcover/run_msi.py new file mode 100644 index 00000000..f4978ade --- /dev/null +++ b/scripts/worldcover/run_msi.py @@ -0,0 +1,419 @@ +#!/usr/bin/env python3 + +import sys +sys.path.append("../../") + +import os +import tempfile +from math import floor +from pathlib import Path +import requests + +import boto3 +import einops +import geopandas as gpd +import pandas as pd +import numpy +import pyarrow as pa +import rasterio +import shapely +import torch +import xarray as xr +from rasterio.windows import Window +from shapely import box +from torchvision.transforms import v2 + +from src.datamodule import ClayDataset +from src.model_clay import CLAYModule + +YEAR = int(os.environ.get("YEAR", 2020)) +DATE = f"{YEAR}-06-01" +TILE_SIZE = 12000 +PATCH_SIZE = 32 +CHIP_SIZE = 512 +E_W_INDEX_START = 67 +E_W_INDEX_END = 125 +N_S_INDEX_START = 24 +N_S_INDEX_END = 49 +YORIGIN = 50.0 +XORIGIN = -125.0 +PXSIZE = 8.333333333333333e-05 + +RASTER_X_SIZE = (E_W_INDEX_END - E_W_INDEX_START) * TILE_SIZE +RASTER_Y_SIZE = (N_S_INDEX_END - N_S_INDEX_START) * TILE_SIZE +NODATA = 0 +CKPT_PATH = ( + "https://huggingface.co/made-with-clay/Clay/resolve/main/" + "Clay_v0.1_epoch-24_val-loss-0.46.ckpt" +) +VERSION = "006" +BUCKET = "clay-worldcover-embeddings" +URL_RGBNIR = "https://esa-worldcover-s2.s3.amazonaws.com/rgbnir/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}W{xidx}_S2RGBNIR.tif" +URL_SWIR = "https://esa-worldcover-s2.s3.amazonaws.com/swir/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}W{xidx}_SWIR.tif" +WC_VERSION_LOOKUP = { + 2020: 100, + 2021: 200, +} + +# Mean and standard deviation for RGBNIR bands +MEAN_RGBNIR = [1369.03, 1597.68, 1741.10, 2858.43] +STD_RGBNIR = [2026.96, 2011.88, 2146.35, 2016.38] + +# Mean and standard deviation for SWIR bands +MEAN_SWIR = [2303.00, 1807.79] +STD_SWIR = [1679.88, 1568.06] + +grid = gpd.read_file( + "https://clay-mgrs-samples.s3.amazonaws.com/esa_worldcover_grid_usa.fgb" +) + + +def tiles_and_windows(input: Window): + print("Input", input) + x_tile_index = E_W_INDEX_END - floor(input.col_off / TILE_SIZE) + x_local_off = input.col_off % TILE_SIZE + x_size = min(CHIP_SIZE, TILE_SIZE - x_local_off) + x_another = x_size < CHIP_SIZE + + y_tile_index = N_S_INDEX_END - floor(input.row_off / TILE_SIZE) + y_local_off = input.row_off % TILE_SIZE + y_size = min(CHIP_SIZE, TILE_SIZE - y_local_off) + y_another = y_size < CHIP_SIZE + + tile_id = f"N{y_tile_index}W{str(x_tile_index).zfill(3)}" + if tile_id not in grid.tile.values: + return + + result = [ + ( + URL_RGBNIR.format( + yidx=y_tile_index, + xidx=str(x_tile_index).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + URL_SWIR.format( + yidx=y_tile_index, + xidx=str(x_tile_index).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + Window(x_local_off, y_local_off, x_size, y_size), + ) + ] + + if x_another: + result.append( + ( + URL_RGBNIR.format( + yidx=y_tile_index, + xidx=str(x_tile_index - 1).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + URL_SWIR.format( + yidx=y_tile_index, + xidx=str(x_tile_index - 1).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + Window(0, y_local_off, CHIP_SIZE - x_size, y_size), + ) + ) + if y_another: + result.append( + ( + URL_RGBNIR.format( + yidx=y_tile_index - 1, + xidx=str(x_tile_index).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + URL_SWIR.format( + yidx=y_tile_index - 1, + xidx=str(x_tile_index).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + Window(x_local_off, 0, x_size, CHIP_SIZE - y_size), + ) + ) + if x_another and y_another: + result.append( + ( + URL_RGBNIR.format( + yidx=y_tile_index - 1, + xidx=str(x_tile_index - 1).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + URL_SWIR.format( + yidx=y_tile_index - 1, + xidx=str(x_tile_index - 1).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), + Window(0, 0, CHIP_SIZE - x_size, CHIP_SIZE - y_size), + ) + ) + + return result + +def download_image(url): + # Download the image from the URL + response = requests.get(url) + # Check if the request was successful + if response.status_code == 200: + return response.content # Return the image content + else: + raise Exception("Failed to download the image") + +def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): + # Download the image from the URL + image_data = download_image(url) + + # Open the image using rasterio from memory + with rasterio.io.MemoryFile(image_data) as memfile: + with memfile.open() as src: + # Read the image data and metadata + img_data = src.read() + img_meta = src.profile + img_crs = src.crs + + # Convert raster data and metadata into an xarray DataArray + img_da = xr.DataArray(img_data, dims=("band", "y", "x"), attrs=img_meta) + + # Tile the data + ds_chunked = img_da.chunk({"y": chunk_size[0], "x": chunk_size[1]}) + + # Get the geospatial information from the original dataset + transform = img_meta["transform"] + + # Iterate over the chunks and compute the geospatial bounds for each chunk + chunk_bounds = {} + + for x in range(ds_chunked.sizes["x"] // chunk_size[1]): + for y in range(ds_chunked.sizes["y"] // chunk_size[0]): + # Compute chunk coordinates + x_start = x * chunk_size[1] + y_start = y * chunk_size[0] + x_end = min(x_start + chunk_size[1], ds_chunked.sizes["x"]) + y_end = min(y_start + chunk_size[0], ds_chunked.sizes["y"]) + + # Compute chunk geospatial bounds + lon_start, lat_start = transform * (x_start, y_start) + lon_end, lat_end = transform * (x_end, y_end) + + # Store chunk bounds + chunk_bounds[(x, y)] = { + "lon_start": lon_start, + "lat_start": lat_start, + "lon_end": lon_end, + "lat_end": lat_end, + } + + return chunk_bounds, img_crs + +def make_batch(result): + rgb_bands = [] + swir_bands = [] + + for url_rgb, url_swir, win in result: + with rasterio.open(url_rgb) as src_rgb, rasterio.open(url_swir) as src_swir: + data_rgb = src_rgb.read(window=win) + data_swir = src_swir.read(window=win) + if NODATA in data_rgb or NODATA in data_swir: + return + transform = src_rgb.window_transform(win) + rgb_bands.append(data_rgb) + swir_bands.append(data_swir) + + if len(rgb_bands) == 0 or len(swir_bands) == 0: + return + + rgb_data = numpy.vstack(rgb_bands) + #rgb_data = rgb_data.transpose(1,2,0) + swir_data = numpy.vstack(swir_bands) + #swir_data = swir_data.transpose(1,2,0) + print("rgb_data: ", rgb_data.shape) + print("swir_data: ", swir_data.shape) + + if rgb_data.shape[0] == 1: + rgb_data = rgb_data[0] + elif rgb_data.shape[0] == 2: + if rgb_data.shape[2] == CHIP_SIZE: + rgb_data = einops.pack(rgb_data, "b * w")[0] + swir_data = einops.pack(swir_data, "b * w")[0] + print("swir_data r1: ", swir_data.shape) + else: + rgb_data = einops.pack(rgb_data, "b h *")[0] + swir_data = einops.pack(swir_data, "b h *")[0] + print("swir_data r2: ", swir_data.shape) + else: + rgb_px1 = einops.pack(rgb_data[:2], "b w *")[0] + rgb_px2 = einops.pack(rgb_data[2:], "b w *")[0] + #rgb_data = einops.pack((rgb_px1, rgb_px2), "b * w")[0] + print("rgb_data re: ", rgb_data.shape) + + #swir_px1 = einops.pack(swir_data[:2], "b w *")[0] + #print("swir_data re: ", swir_px1.shape) + #swir_px2 = einops.pack(swir_data[2:], "b w *")[0] + #swir_data = einops.pack((swir_px1, swir_px2), "b * w")[0] + print("swir_data re: ", swir_data.shape) + + rgb_data = rgb_data.transpose(1,2,0) + swir_data = swir_data.transpose(1,2,0) + print("rgb_data: ", rgb_data.shape) + print("swir_data: ", swir_data.shape) + combined_data = numpy.concatenate((rgb_data,swir_data), axis=-1) #numpy.concatenate([rgb_data, swir_data]) + combined_data = combined_data.transpose(2,0,1) + print("combined_data: ", combined_data.shape) + + return { + "pixels": torch.as_tensor(data=[combined_data], dtype=torch.float32).to(rgb_model.device), + "latlon": torch.as_tensor(data=[ds.normalize_latlon(transform[0], transform[3])]).to(rgb_model.device), + "timestep": torch.as_tensor(data=[ds.normalize_timestamp(f"{YEAR}-06-01")]).to(rgb_model.device), + "date": f"{YEAR}-06-01" + } + + + +index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 2)) + +# Setup model components +tfm = v2.Compose([v2.Normalize(mean=MEAN_RGBNIR + MEAN_SWIR, std=STD_RGBNIR + STD_SWIR)]) +ds = ClayDataset(chips_path=[], transform=tfm) + +# Load model +rgb_model = CLAYModule.load_from_checkpoint( + CKPT_PATH, + mask_ratio=0.0, + band_groups={"rgb": (2, 1, 0), "nir": (3,), "swir": (4, 5)}, + bands=6, + strict=False, # ignore the extra parameters in the checkpoint + embeddings_level="group", +) +# Set the model to evaluation mode +rgb_model.eval() + +outdir_embeddings = Path("data/embeddings") +outdir_embeddings.mkdir(exist_ok=True, parents=True) + +xoff = index * CHIP_SIZE +yoff = 0 +embeddings = [] +all_bounds = [] +results = [] +while yoff < RASTER_Y_SIZE: + result = tiles_and_windows(Window(xoff, yoff, CHIP_SIZE, CHIP_SIZE)) + if result is not None: + results.append(result) + + if result is None: + yoff += CHIP_SIZE + continue + + batch = make_batch(result) + if batch is None: + yoff += CHIP_SIZE + continue + + ( + unmasked_patches, + unmasked_indices, + masked_indices, + masked_matrix, + ) = rgb_model.model.encoder(batch) + + embeddings.append(unmasked_patches.detach().cpu().numpy()) + all_bounds.append( + ( + XORIGIN + PXSIZE * xoff, + YORIGIN - PXSIZE * (yoff + CHIP_SIZE), + XORIGIN + PXSIZE * (xoff + CHIP_SIZE), + YORIGIN - PXSIZE * yoff, + ) + ) + + yoff += CHIP_SIZE + + print(len(embeddings), len(results)) + embeddings_ = numpy.vstack(embeddings) + #embeddings_ = embeddings[0] + print("Embeddings shape: ", embeddings_.shape) + + # remove date and lat/lon + embeddings_ = embeddings_[:, :-2, :].mean(axis=0) + + print(f"Embeddings have shape {embeddings_.shape}") + + # reshape to disaggregated patches + embeddings_patch = embeddings_.reshape([3, 16, 16, 768]) + + # average over the band groups + embeddings_mean = embeddings_patch.mean(axis=0) + + print(f"Average patch embeddings have shape {embeddings_mean.shape}") + + if result is not None: + print("result: ", result[0][0]) + #pix = get_pixels(result) + chunk_bounds, epsg = patch_bounds_from_url(result[0][0]) + #print("chunk_bounds: ", chunk_bounds) + print("chunk bounds length:", len(chunk_bounds)) + + + # Iterate through each patch + for i in range(embeddings_mean.shape[0]): + for j in range(embeddings_mean.shape[1]): + embeddings_output_patch = embeddings_mean[i, j] + + item_ = [ + element for element in list(chunk_bounds.items()) if element[0] == (i, j) + ] + box_ = [ + item_[0][1]["lon_start"], + item_[0][1]["lat_start"], + item_[0][1]["lon_end"], + item_[0][1]["lat_end"], + ] + + data = { + #"source_url": batch["source_url"][0], + #"date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( + # dtype="date32[day][pyarrow]" + #), + #"date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), + "date": pd.to_datetime(batch["date"], format="%Y-%m-%d"), + "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], + } + + # Define the bounding box as a Polygon (xmin, ymin, xmax, ymax) + # The box_ list is encoded as + # [bottom left x, bottom left y, top right x, top right y] + box_emb = shapely.geometry.box(box_[0], box_[1], box_[2], box_[3]) + + print(str(epsg)[-4:]) + + # Create the GeoDataFrame + gdf = gpd.GeoDataFrame(data, geometry=[box_emb], crs=f"EPSG:{str(epsg)[-4:]}") + + # Reproject to WGS84 (lon/lat coordinates) + gdf = gdf.to_crs(epsg=4326) + + with tempfile.TemporaryDirectory() as tmp: + # tmp = "/home/tam/Desktop/wcctmp" + + outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_v{VERSION}.gpq" + print(f"Uploading embeddings to {outpath}") + #print(gdf) + + gdf.to_parquet(path=outpath, compression="ZSTD", schema_version="1.0.0") + + s3_client = boto3.client("s3") + s3_client.upload_file( + outpath, + BUCKET, + f"v{VERSION}/{YEAR}/{os.path.basename(outpath)}", + ) + From be3cc333bb6ecc863bb65a257efe171875f26211 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 1 May 2024 14:53:48 -0700 Subject: [PATCH 06/11] incorporate sar composite --- .../worldcover/{run_msi.py => run_msi_sar.py} | 97 +++++++++---------- 1 file changed, 47 insertions(+), 50 deletions(-) rename scripts/worldcover/{run_msi.py => run_msi_sar.py} (85%) diff --git a/scripts/worldcover/run_msi.py b/scripts/worldcover/run_msi_sar.py similarity index 85% rename from scripts/worldcover/run_msi.py rename to scripts/worldcover/run_msi_sar.py index f4978ade..e84d6cf2 100644 --- a/scripts/worldcover/run_msi.py +++ b/scripts/worldcover/run_msi_sar.py @@ -46,10 +46,11 @@ "https://huggingface.co/made-with-clay/Clay/resolve/main/" "Clay_v0.1_epoch-24_val-loss-0.46.ckpt" ) -VERSION = "006" +VERSION = "007" BUCKET = "clay-worldcover-embeddings" URL_RGBNIR = "https://esa-worldcover-s2.s3.amazonaws.com/rgbnir/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}W{xidx}_S2RGBNIR.tif" URL_SWIR = "https://esa-worldcover-s2.s3.amazonaws.com/swir/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}W{xidx}_SWIR.tif" +URL_SAR = "https://esa-worldcover-s1.s3.amazonaws.com/vvvhratio/{year}/N{yidx}/ESA_WorldCover_10m_{year}_v{version}_N{yidx}E{xidx}_S1VVVHratio.tif" WC_VERSION_LOOKUP = { 2020: 100, 2021: 200, @@ -63,6 +64,10 @@ MEAN_SWIR = [2303.00, 1807.79] STD_SWIR = [1679.88, 1568.06] +# Mean and standard deviation for SAR bands +MEAN_SAR = [0.026, 0.118, 0.118] +STD_SAR = [0.118, 0.873, 0.873] + grid = gpd.read_file( "https://clay-mgrs-samples.s3.amazonaws.com/esa_worldcover_grid_usa.fgb" ) @@ -98,6 +103,12 @@ def tiles_and_windows(input: Window): year=YEAR, version=WC_VERSION_LOOKUP[YEAR], ), + URL_SAR.format( + yidx=y_tile_index, + xidx=str(x_tile_index).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), Window(x_local_off, y_local_off, x_size, y_size), ) ] @@ -117,6 +128,12 @@ def tiles_and_windows(input: Window): year=YEAR, version=WC_VERSION_LOOKUP[YEAR], ), + URL_SAR.format( + yidx=y_tile_index, + xidx=str(x_tile_index - 1).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), Window(0, y_local_off, CHIP_SIZE - x_size, y_size), ) ) @@ -135,6 +152,12 @@ def tiles_and_windows(input: Window): year=YEAR, version=WC_VERSION_LOOKUP[YEAR], ), + URL_SAR.format( + yidx=y_tile_index - 1, + xidx=str(x_tile_index).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), Window(x_local_off, 0, x_size, CHIP_SIZE - y_size), ) ) @@ -153,6 +176,12 @@ def tiles_and_windows(input: Window): year=YEAR, version=WC_VERSION_LOOKUP[YEAR], ), + URL_SAR.format( + yidx=y_tile_index - 1, + xidx=str(x_tile_index - 1).zfill(3), + year=YEAR, + version=WC_VERSION_LOOKUP[YEAR], + ), Window(0, 0, CHIP_SIZE - x_size, CHIP_SIZE - y_size), ) ) @@ -217,57 +246,31 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): def make_batch(result): rgb_bands = [] swir_bands = [] + sar_bands = [] - for url_rgb, url_swir, win in result: - with rasterio.open(url_rgb) as src_rgb, rasterio.open(url_swir) as src_swir: + for url_rgb, url_swir, url_sar, win in result: + with rasterio.open(url_rgb) as src_rgb, rasterio.open(url_swir) as src_swir, rasterio.open(url_sar) as src_sar: data_rgb = src_rgb.read(window=win) data_swir = src_swir.read(window=win) - if NODATA in data_rgb or NODATA in data_swir: + data_sar = src_sar.read(window=win) + if NODATA in data_rgb or NODATA in data_swir or NODATA in data_sar: return transform = src_rgb.window_transform(win) rgb_bands.append(data_rgb) swir_bands.append(data_swir) + sar_bands.append(data_sar) - if len(rgb_bands) == 0 or len(swir_bands) == 0: + if len(rgb_bands) == 0 or len(swir_bands) == 0 or len(sar_bands) == 0: return rgb_data = numpy.vstack(rgb_bands) - #rgb_data = rgb_data.transpose(1,2,0) swir_data = numpy.vstack(swir_bands) - #swir_data = swir_data.transpose(1,2,0) - print("rgb_data: ", rgb_data.shape) - print("swir_data: ", swir_data.shape) - - if rgb_data.shape[0] == 1: - rgb_data = rgb_data[0] - elif rgb_data.shape[0] == 2: - if rgb_data.shape[2] == CHIP_SIZE: - rgb_data = einops.pack(rgb_data, "b * w")[0] - swir_data = einops.pack(swir_data, "b * w")[0] - print("swir_data r1: ", swir_data.shape) - else: - rgb_data = einops.pack(rgb_data, "b h *")[0] - swir_data = einops.pack(swir_data, "b h *")[0] - print("swir_data r2: ", swir_data.shape) - else: - rgb_px1 = einops.pack(rgb_data[:2], "b w *")[0] - rgb_px2 = einops.pack(rgb_data[2:], "b w *")[0] - #rgb_data = einops.pack((rgb_px1, rgb_px2), "b * w")[0] - print("rgb_data re: ", rgb_data.shape) - - #swir_px1 = einops.pack(swir_data[:2], "b w *")[0] - #print("swir_data re: ", swir_px1.shape) - #swir_px2 = einops.pack(swir_data[2:], "b w *")[0] - #swir_data = einops.pack((swir_px1, swir_px2), "b * w")[0] - print("swir_data re: ", swir_data.shape) - - rgb_data = rgb_data.transpose(1,2,0) - swir_data = swir_data.transpose(1,2,0) - print("rgb_data: ", rgb_data.shape) - print("swir_data: ", swir_data.shape) - combined_data = numpy.concatenate((rgb_data,swir_data), axis=-1) #numpy.concatenate([rgb_data, swir_data]) - combined_data = combined_data.transpose(2,0,1) - print("combined_data: ", combined_data.shape) + sar_data = numpy.vstack(sar_bands) + + # Normalize SAR data + #sar_data = (sar_data - MEAN_SAR) / STD_SAR + + combined_data = numpy.concatenate((rgb_data, swir_data, sar_data), axis=0) return { "pixels": torch.as_tensor(data=[combined_data], dtype=torch.float32).to(rgb_model.device), @@ -276,20 +279,18 @@ def make_batch(result): "date": f"{YEAR}-06-01" } - - index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 2)) # Setup model components -tfm = v2.Compose([v2.Normalize(mean=MEAN_RGBNIR + MEAN_SWIR, std=STD_RGBNIR + STD_SWIR)]) +tfm = v2.Compose([v2.Normalize(mean=MEAN_RGBNIR + MEAN_SWIR + MEAN_SAR, std=STD_RGBNIR + STD_SWIR + STD_SAR)]) ds = ClayDataset(chips_path=[], transform=tfm) # Load model rgb_model = CLAYModule.load_from_checkpoint( CKPT_PATH, mask_ratio=0.0, - band_groups={"rgb": (2, 1, 0), "nir": (3,), "swir": (4, 5)}, - bands=6, + band_groups={"rgb": (2, 1, 0), "nir": (3,), "swir": (4, 5), "sar": (6, 7)}, + bands=8, strict=False, # ignore the extra parameters in the checkpoint embeddings_level="group", ) @@ -339,7 +340,6 @@ def make_batch(result): print(len(embeddings), len(results)) embeddings_ = numpy.vstack(embeddings) - #embeddings_ = embeddings[0] print("Embeddings shape: ", embeddings_.shape) # remove date and lat/lon @@ -348,7 +348,7 @@ def make_batch(result): print(f"Embeddings have shape {embeddings_.shape}") # reshape to disaggregated patches - embeddings_patch = embeddings_.reshape([3, 16, 16, 768]) + embeddings_patch = embeddings_.reshape([4, 16, 16, 768]) # average over the band groups embeddings_mean = embeddings_patch.mean(axis=0) @@ -357,9 +357,7 @@ def make_batch(result): if result is not None: print("result: ", result[0][0]) - #pix = get_pixels(result) chunk_bounds, epsg = patch_bounds_from_url(result[0][0]) - #print("chunk_bounds: ", chunk_bounds) print("chunk bounds length:", len(chunk_bounds)) @@ -416,4 +414,3 @@ def make_batch(result): BUCKET, f"v{VERSION}/{YEAR}/{os.path.basename(outpath)}", ) - From f17ae75536c27ffc171efcc503ade31111acd26b Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 1 May 2024 15:58:40 -0700 Subject: [PATCH 07/11] incorporate sar composite --- scripts/worldcover/run_msi_sar.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/worldcover/run_msi_sar.py b/scripts/worldcover/run_msi_sar.py index e84d6cf2..075af20f 100644 --- a/scripts/worldcover/run_msi_sar.py +++ b/scripts/worldcover/run_msi_sar.py @@ -65,8 +65,8 @@ STD_SWIR = [1679.88, 1568.06] # Mean and standard deviation for SAR bands -MEAN_SAR = [0.026, 0.118, 0.118] -STD_SAR = [0.118, 0.873, 0.873] +MEAN_SAR = [0.026, 0.118] +STD_SAR = [0.118, 0.873] grid = gpd.read_file( "https://clay-mgrs-samples.s3.amazonaws.com/esa_worldcover_grid_usa.fgb" From af0a3937bca7a8601da8fd7b26091e9f24aa7a79 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 May 2024 22:59:30 +0000 Subject: [PATCH 08/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/worldcover/embeddings_db.py | 4 +- scripts/worldcover/run.py | 89 +++++++++++++------------- scripts/worldcover/run_msi_sar.py | 97 ++++++++++++++++++----------- 3 files changed, 107 insertions(+), 83 deletions(-) diff --git a/scripts/worldcover/embeddings_db.py b/scripts/worldcover/embeddings_db.py index bb53ae02..b417dbf4 100644 --- a/scripts/worldcover/embeddings_db.py +++ b/scripts/worldcover/embeddings_db.py @@ -31,7 +31,7 @@ db.table_names() # Drop existing table if exists -#db.drop_table("worldcover-2020-v001") +# db.drop_table("worldcover-2020-v001") # Create embeddings table and insert the vector data tbl = db.create_table("worldcover-2020-v001", data=data, mode="overwrite") @@ -55,4 +55,4 @@ def plot(df, cols=10): # Select a vector by index, and search 10 similar pairs, and plot v = tbl.to_pandas()["vector"].values[5] result = tbl.search(query=v).limit(5).to_pandas() -plot(result, 5) \ No newline at end of file +plot(result, 5) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index 95ad71b5..62b8f5ed 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -1,26 +1,25 @@ #!/usr/bin/env python3 import sys + sys.path.append("../../") import os import tempfile from math import floor from pathlib import Path -import requests import boto3 import einops import geopandas as gpd -import pandas as pd import numpy -import pyarrow as pa +import pandas as pd import rasterio +import requests import shapely import torch import xarray as xr from rasterio.windows import Window -from shapely import box from torchvision.transforms import v2 from src.datamodule import ClayDataset @@ -141,6 +140,7 @@ def tiles_and_windows(input: Window): return result + def download_image(url): # Download an image from a URL response = requests.get(url) @@ -150,10 +150,11 @@ def download_image(url): else: raise Exception("Failed to download the image") + def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): # Download an image from a URL image_data = download_image(url) - + # Open the image using rasterio from memory with rasterio.io.MemoryFile(image_data) as memfile: with memfile.open() as src: @@ -161,19 +162,19 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): img_data = src.read() img_meta = src.profile img_crs = src.crs - + # Convert raster data and metadata into an xarray DataArray img_da = xr.DataArray(img_data, dims=("band", "y", "x"), attrs=img_meta) - + # Tile the data ds_chunked = img_da.chunk({"y": chunk_size[0], "x": chunk_size[1]}) - + # Get the geospatial information from the original dataset transform = img_meta["transform"] - + # Iterate over the chunks and compute the geospatial bounds for each chunk chunk_bounds = {} - + for x in range(ds_chunked.sizes["x"] // chunk_size[1]): for y in range(ds_chunked.sizes["y"] // chunk_size[0]): # Compute chunk coordinates @@ -181,11 +182,11 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): y_start = y * chunk_size[0] x_end = min(x_start + chunk_size[1], ds_chunked.sizes["x"]) y_end = min(y_start + chunk_size[0], ds_chunked.sizes["y"]) - + # Compute chunk geospatial bounds lon_start, lat_start = transform * (x_start, y_start) lon_end, lat_end = transform * (x_end, y_end) - + # Store chunk bounds chunk_bounds[(x, y)] = { "lon_start": lon_start, @@ -193,9 +194,10 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): "lon_end": lon_end, "lat_end": lat_end, } - + return chunk_bounds, img_crs + def make_batch(result): pixels = [] for url, win in result: @@ -230,10 +232,10 @@ def make_batch(result): "timestep": torch.as_tensor(data=[ds.normalize_timestamp(f"{YEAR}-06-01")]).to( rgb_model.device ), - "date": f"{YEAR}-06-01" - , + "date": f"{YEAR}-06-01", } + def get_pixels(result): pixels = [] for url, win in result: @@ -319,42 +321,41 @@ def get_pixels(result): ) yoff += CHIP_SIZE - - print(len(embeddings), len(results)) embeddings_ = numpy.vstack(embeddings) - #embeddings_ = embeddings[0] + # embeddings_ = embeddings[0] print("Embeddings shape: ", embeddings_.shape) # remove date and lat/lon embeddings_ = embeddings_[:, :-2, :].mean(axis=0) - + print(f"Embeddings have shape {embeddings_.shape}") - + # reshape to disaggregated patches embeddings_patch = embeddings_.reshape([2, 16, 16, 768]) - + # average over the band groups embeddings_mean = embeddings_patch.mean(axis=0) - - print(f"Average patch embeddings have shape {embeddings_mean.shape}") + print(f"Average patch embeddings have shape {embeddings_mean.shape}") if result is not None: print("result: ", result[0][0]) pix = get_pixels(result) chunk_bounds, epsg = patch_bounds_from_url(result[0][0]) - #print("chunk_bounds: ", chunk_bounds) + # print("chunk_bounds: ", chunk_bounds) print("chunk bounds length:", len(chunk_bounds)) - + # Iterate through each patch for i in range(embeddings_mean.shape[0]): for j in range(embeddings_mean.shape[1]): embeddings_output_patch = embeddings_mean[i, j] - + item_ = [ - element for element in list(chunk_bounds.items()) if element[0] == (i, j) + element + for element in list(chunk_bounds.items()) + if element[0] == (i, j) ] box_ = [ item_[0][1]["lon_start"], @@ -364,42 +365,44 @@ def get_pixels(result): ] data = { - #"source_url": batch["source_url"][0], - #"date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( + # "source_url": batch["source_url"][0], + # "date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( # dtype="date32[day][pyarrow]" - #), - #"date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), + # ), + # "date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), "date": pd.to_datetime(batch["date"], format="%Y-%m-%d"), "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], } - + # Define the bounding box as a Polygon (xmin, ymin, xmax, ymax) # The box_ list is encoded as # [bottom left x, bottom left y, top right x, top right y] box_emb = shapely.geometry.box(box_[0], box_[1], box_[2], box_[3]) print(str(epsg)[-4:]) - + # Create the GeoDataFrame - gdf = gpd.GeoDataFrame(data, geometry=[box_emb], crs=f"EPSG:{str(epsg)[-4:]}") - + gdf = gpd.GeoDataFrame( + data, geometry=[box_emb], crs=f"EPSG:{str(epsg)[-4:]}" + ) + # Reproject to WGS84 (lon/lat coordinates) gdf = gdf.to_crs(epsg=4326) - + with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" - + outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_v{VERSION}.gpq" print(f"Uploading embeddings to {outpath}") - #print(gdf) - - gdf.to_parquet(path=outpath, compression="ZSTD", schema_version="1.0.0") - + # print(gdf) + + gdf.to_parquet( + path=outpath, compression="ZSTD", schema_version="1.0.0" + ) + s3_client = boto3.client("s3") s3_client.upload_file( outpath, BUCKET, f"v{VERSION}/{YEAR}/{os.path.basename(outpath)}", ) - - \ No newline at end of file diff --git a/scripts/worldcover/run_msi_sar.py b/scripts/worldcover/run_msi_sar.py index 075af20f..b8775bcd 100644 --- a/scripts/worldcover/run_msi_sar.py +++ b/scripts/worldcover/run_msi_sar.py @@ -1,26 +1,24 @@ #!/usr/bin/env python3 import sys + sys.path.append("../../") import os import tempfile from math import floor from pathlib import Path -import requests import boto3 -import einops import geopandas as gpd -import pandas as pd import numpy -import pyarrow as pa +import pandas as pd import rasterio +import requests import shapely import torch import xarray as xr from rasterio.windows import Window -from shapely import box from torchvision.transforms import v2 from src.datamodule import ClayDataset @@ -188,6 +186,7 @@ def tiles_and_windows(input: Window): return result + def download_image(url): # Download the image from the URL response = requests.get(url) @@ -197,10 +196,11 @@ def download_image(url): else: raise Exception("Failed to download the image") + def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): # Download the image from the URL image_data = download_image(url) - + # Open the image using rasterio from memory with rasterio.io.MemoryFile(image_data) as memfile: with memfile.open() as src: @@ -208,19 +208,19 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): img_data = src.read() img_meta = src.profile img_crs = src.crs - + # Convert raster data and metadata into an xarray DataArray img_da = xr.DataArray(img_data, dims=("band", "y", "x"), attrs=img_meta) - + # Tile the data ds_chunked = img_da.chunk({"y": chunk_size[0], "x": chunk_size[1]}) - + # Get the geospatial information from the original dataset transform = img_meta["transform"] - + # Iterate over the chunks and compute the geospatial bounds for each chunk chunk_bounds = {} - + for x in range(ds_chunked.sizes["x"] // chunk_size[1]): for y in range(ds_chunked.sizes["y"] // chunk_size[0]): # Compute chunk coordinates @@ -228,11 +228,11 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): y_start = y * chunk_size[0] x_end = min(x_start + chunk_size[1], ds_chunked.sizes["x"]) y_end = min(y_start + chunk_size[0], ds_chunked.sizes["y"]) - + # Compute chunk geospatial bounds lon_start, lat_start = transform * (x_start, y_start) lon_end, lat_end = transform * (x_end, y_end) - + # Store chunk bounds chunk_bounds[(x, y)] = { "lon_start": lon_start, @@ -240,16 +240,19 @@ def patch_bounds_from_url(url, chunk_size=(PATCH_SIZE, PATCH_SIZE)): "lon_end": lon_end, "lat_end": lat_end, } - + return chunk_bounds, img_crs + def make_batch(result): rgb_bands = [] swir_bands = [] sar_bands = [] for url_rgb, url_swir, url_sar, win in result: - with rasterio.open(url_rgb) as src_rgb, rasterio.open(url_swir) as src_swir, rasterio.open(url_sar) as src_sar: + with rasterio.open(url_rgb) as src_rgb, rasterio.open( + url_swir + ) as src_swir, rasterio.open(url_sar) as src_sar: data_rgb = src_rgb.read(window=win) data_swir = src_swir.read(window=win) data_sar = src_sar.read(window=win) @@ -268,21 +271,34 @@ def make_batch(result): sar_data = numpy.vstack(sar_bands) # Normalize SAR data - #sar_data = (sar_data - MEAN_SAR) / STD_SAR + # sar_data = (sar_data - MEAN_SAR) / STD_SAR combined_data = numpy.concatenate((rgb_data, swir_data, sar_data), axis=0) return { - "pixels": torch.as_tensor(data=[combined_data], dtype=torch.float32).to(rgb_model.device), - "latlon": torch.as_tensor(data=[ds.normalize_latlon(transform[0], transform[3])]).to(rgb_model.device), - "timestep": torch.as_tensor(data=[ds.normalize_timestamp(f"{YEAR}-06-01")]).to(rgb_model.device), - "date": f"{YEAR}-06-01" + "pixels": torch.as_tensor(data=[combined_data], dtype=torch.float32).to( + rgb_model.device + ), + "latlon": torch.as_tensor( + data=[ds.normalize_latlon(transform[0], transform[3])] + ).to(rgb_model.device), + "timestep": torch.as_tensor(data=[ds.normalize_timestamp(f"{YEAR}-06-01")]).to( + rgb_model.device + ), + "date": f"{YEAR}-06-01", } + index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 2)) # Setup model components -tfm = v2.Compose([v2.Normalize(mean=MEAN_RGBNIR + MEAN_SWIR + MEAN_SAR, std=STD_RGBNIR + STD_SWIR + STD_SAR)]) +tfm = v2.Compose( + [ + v2.Normalize( + mean=MEAN_RGBNIR + MEAN_SWIR + MEAN_SAR, std=STD_RGBNIR + STD_SWIR + STD_SAR + ) + ] +) ds = ClayDataset(chips_path=[], transform=tfm) # Load model @@ -360,14 +376,15 @@ def make_batch(result): chunk_bounds, epsg = patch_bounds_from_url(result[0][0]) print("chunk bounds length:", len(chunk_bounds)) - # Iterate through each patch for i in range(embeddings_mean.shape[0]): for j in range(embeddings_mean.shape[1]): embeddings_output_patch = embeddings_mean[i, j] - + item_ = [ - element for element in list(chunk_bounds.items()) if element[0] == (i, j) + element + for element in list(chunk_bounds.items()) + if element[0] == (i, j) ] box_ = [ item_[0][1]["lon_start"], @@ -377,37 +394,41 @@ def make_batch(result): ] data = { - #"source_url": batch["source_url"][0], - #"date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( + # "source_url": batch["source_url"][0], + # "date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( # dtype="date32[day][pyarrow]" - #), - #"date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), + # ), + # "date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), "date": pd.to_datetime(batch["date"], format="%Y-%m-%d"), "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], } - + # Define the bounding box as a Polygon (xmin, ymin, xmax, ymax) # The box_ list is encoded as # [bottom left x, bottom left y, top right x, top right y] box_emb = shapely.geometry.box(box_[0], box_[1], box_[2], box_[3]) print(str(epsg)[-4:]) - + # Create the GeoDataFrame - gdf = gpd.GeoDataFrame(data, geometry=[box_emb], crs=f"EPSG:{str(epsg)[-4:]}") - + gdf = gpd.GeoDataFrame( + data, geometry=[box_emb], crs=f"EPSG:{str(epsg)[-4:]}" + ) + # Reproject to WGS84 (lon/lat coordinates) gdf = gdf.to_crs(epsg=4326) - + with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" - + outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_v{VERSION}.gpq" print(f"Uploading embeddings to {outpath}") - #print(gdf) - - gdf.to_parquet(path=outpath, compression="ZSTD", schema_version="1.0.0") - + # print(gdf) + + gdf.to_parquet( + path=outpath, compression="ZSTD", schema_version="1.0.0" + ) + s3_client = boto3.client("s3") s3_client.upload_file( outpath, From 52e302eeae0fb231fdfd167be307f46d85ab1ea6 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Tue, 21 May 2024 19:11:39 -0400 Subject: [PATCH 09/11] fixes for linting --- scripts/worldcover/run.py | 14 ++++++-------- scripts/worldcover/run_msi_sar.py | 12 +++++------- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index 62b8f5ed..5a720a3b 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -37,6 +37,7 @@ YORIGIN = 50.0 XORIGIN = -125.0 PXSIZE = 8.333333333333333e-05 +SUCCESS_CODE = 200 RASTER_X_SIZE = (E_W_INDEX_END - E_W_INDEX_START) * TILE_SIZE RASTER_Y_SIZE = (N_S_INDEX_END - N_S_INDEX_START) * TILE_SIZE @@ -145,7 +146,7 @@ def download_image(url): # Download an image from a URL response = requests.get(url) # Check if the request was successful - if response.status_code == 200: + if response.status_code == SUCCESS_CODE: return response.content # Return the image content else: raise Exception("Failed to download the image") @@ -244,7 +245,7 @@ def get_pixels(result): if NODATA in data: return pixels.append(data) - transform = src.window_transform(win) + # transform = src.window_transform(win) if len(pixels) == 1: pixels = pixels[0] @@ -365,11 +366,6 @@ def get_pixels(result): ] data = { - # "source_url": batch["source_url"][0], - # "date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( - # dtype="date32[day][pyarrow]" - # ), - # "date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), "date": pd.to_datetime(batch["date"], format="%Y-%m-%d"), "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], } @@ -397,7 +393,9 @@ def get_pixels(result): # print(gdf) gdf.to_parquet( - path=outpath, compression="ZSTD", schema_version="1.0.0" + path=outpath, + compression="ZSTD", + schema_version="1.0.0" ) s3_client = boto3.client("s3") diff --git a/scripts/worldcover/run_msi_sar.py b/scripts/worldcover/run_msi_sar.py index b8775bcd..fc675bce 100644 --- a/scripts/worldcover/run_msi_sar.py +++ b/scripts/worldcover/run_msi_sar.py @@ -36,6 +36,7 @@ YORIGIN = 50.0 XORIGIN = -125.0 PXSIZE = 8.333333333333333e-05 +SUCCESS_CODE = 200 RASTER_X_SIZE = (E_W_INDEX_END - E_W_INDEX_START) * TILE_SIZE RASTER_Y_SIZE = (N_S_INDEX_END - N_S_INDEX_START) * TILE_SIZE @@ -191,7 +192,7 @@ def download_image(url): # Download the image from the URL response = requests.get(url) # Check if the request was successful - if response.status_code == 200: + if response.status_code == SUCCESS_CODE: return response.content # Return the image content else: raise Exception("Failed to download the image") @@ -394,11 +395,6 @@ def make_batch(result): ] data = { - # "source_url": batch["source_url"][0], - # "date": pd.to_datetime(arg=date, format="%Y-%m-%d").astype( - # dtype="date32[day][pyarrow]" - # ), - # "date": pd.to_datetime(date, format="%Y-%m-%d", dtype="date32[day][pyarrow]"), "date": pd.to_datetime(batch["date"], format="%Y-%m-%d"), "embeddings": [numpy.ascontiguousarray(embeddings_output_patch)], } @@ -426,7 +422,9 @@ def make_batch(result): # print(gdf) gdf.to_parquet( - path=outpath, compression="ZSTD", schema_version="1.0.0" + path=outpath, + compression="ZSTD", + schema_version="1.0.0" ) s3_client = boto3.client("s3") From 48ef01ac5aacf29ee0b0013f629b2cdae2e16710 Mon Sep 17 00:00:00 2001 From: Lilly Thomas Date: Wed, 22 May 2024 12:52:12 -0400 Subject: [PATCH 10/11] fixes for linting --- scripts/worldcover/run.py | 3 ++- scripts/worldcover/run_msi_sar.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index 5a720a3b..3e733fa5 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -388,7 +388,8 @@ def get_pixels(result): with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" - outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_v{VERSION}.gpq" + outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_" \ + f"v{VERSION}.gpq" print(f"Uploading embeddings to {outpath}") # print(gdf) diff --git a/scripts/worldcover/run_msi_sar.py b/scripts/worldcover/run_msi_sar.py index fc675bce..f0c90db5 100644 --- a/scripts/worldcover/run_msi_sar.py +++ b/scripts/worldcover/run_msi_sar.py @@ -417,7 +417,9 @@ def make_batch(result): with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" - outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_v{VERSION}.gpq" + outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_" \ + f"v{VERSION}.gpq" + print(f"Uploading embeddings to {outpath}") # print(gdf) From d4a3c280ff23f8007b1156b08aa6407d2d1f60e0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 May 2024 16:54:11 +0000 Subject: [PATCH 11/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/worldcover/run.py | 10 +++++----- scripts/worldcover/run_msi_sar.py | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/scripts/worldcover/run.py b/scripts/worldcover/run.py index 3e733fa5..69485160 100755 --- a/scripts/worldcover/run.py +++ b/scripts/worldcover/run.py @@ -388,15 +388,15 @@ def get_pixels(result): with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" - outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_" \ - f"v{VERSION}.gpq" + outpath = ( + f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_" + f"v{VERSION}.gpq" + ) print(f"Uploading embeddings to {outpath}") # print(gdf) gdf.to_parquet( - path=outpath, - compression="ZSTD", - schema_version="1.0.0" + path=outpath, compression="ZSTD", schema_version="1.0.0" ) s3_client = boto3.client("s3") diff --git a/scripts/worldcover/run_msi_sar.py b/scripts/worldcover/run_msi_sar.py index f0c90db5..b25d76fd 100644 --- a/scripts/worldcover/run_msi_sar.py +++ b/scripts/worldcover/run_msi_sar.py @@ -417,16 +417,16 @@ def make_batch(result): with tempfile.TemporaryDirectory() as tmp: # tmp = "/home/tam/Desktop/wcctmp" - outpath = f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_" \ - f"v{VERSION}.gpq" + outpath = ( + f"{tmp}/worldcover_patch_embeddings_{YEAR}_{index}_{i}_{j}_" + f"v{VERSION}.gpq" + ) print(f"Uploading embeddings to {outpath}") # print(gdf) gdf.to_parquet( - path=outpath, - compression="ZSTD", - schema_version="1.0.0" + path=outpath, compression="ZSTD", schema_version="1.0.0" ) s3_client = boto3.client("s3")