ai4eo_challenge_template.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "75ac6ced-006a-49d9-95d6-46ac16001ef8",
   "metadata": {},
   "source": [
    "#### Clay Embedding Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f89f38f-641f-419c-bc9d-38ece0f01dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# functions to generate embedding for participant convenience\n",
    "import pystac_client \n",
    "import stackstac\n",
    "import torch\n",
    "import yaml\n",
    "import math\n",
    "import rasterio\n",
    "import warnings\n",
    "import os\n",
    "import joblib\n",
    "\n",
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import numpy as np\n",
    "\n",
    "from shapely.geometry import Point, box\n",
    "from torchvision import transforms as v2\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "from box import Box\n",
    "from rasterio.enums import Resampling\n",
    "from tqdm import tqdm\n",
    "\n",
    "from src.model import ClayMAEModule\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "STAC_API = \"https://earth-search.aws.element84.com/v1\"\n",
    "\n",
    "# Load the model and metadata\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "ckpt = \"https://clay-model-ckpt.s3.amazonaws.com/v0.5.7/mae_v0.5.7_epoch-13_val-loss-0.3098.ckpt\"\n",
    "torch.set_default_device(device)\n",
    "\n",
    "torch.cuda.empty_cache()  # Clear GPU cache\n",
    "\n",
    "model = ClayMAEModule.load_from_checkpoint(\n",
    "    ckpt, metadata_path=\"configs/metadata.yaml\", shuffle=False, mask_ratio=0\n",
    ")\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "\n",
    "metadata = Box(yaml.safe_load(open(\"configs/metadata.yaml\")))\n",
    "\n",
    "# Function to normalize timestamp\n",
    "def normalize_timestamp(date):\n",
    "    week = date.isocalendar().week * 2 * np.pi / 52\n",
    "    hour = date.hour * 2 * np.pi / 24\n",
    "    return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour))\n",
    "\n",
    "# Function to normalize lat/lon\n",
    "def normalize_latlon(lat, lon):\n",
    "    lat = lat * np.pi / 180\n",
    "    lon = lon * np.pi / 180\n",
    "    return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon))\n",
    "\n",
    "def to_device(data, device):\n",
    "    if isinstance(data, torch.Tensor):\n",
    "        return data.to(device)\n",
    "    elif isinstance(data, dict):\n",
    "        return {k: to_device(v, device) for k, v in data.items()}\n",
    "    elif isinstance(data, list):\n",
    "        return [to_device(v, device) for v in data]\n",
    "    return data\n",
    "\n",
    "def process_point(lon, lat, model, metadata, start_date, end_date, COLLECTION, device):\n",
    "    if COLLECTION == \"landsat-c2-l2\":\n",
    "        os.environ['AWS_REQUEST_PAYER'] = 'requester'\n",
    "    model.to(device)  # Ensure the model is on the correct device\n",
    "    catalog = pystac_client.Client.open(STAC_API)\n",
    "    search = catalog.search(\n",
    "        collections=[COLLECTION],\n",
    "        datetime=f\"{start_date}/{end_date}\",\n",
    "        bbox=(lon - 1e-5, lat - 1e-5, lon + 1e-5, lat + 1e-5),\n",
    "        max_items=10,\n",
    "        query={\"eo:cloud_cover\": {\"lt\": 80}},\n",
    "    )\n",
    "\n",
    "    all_items = search.get_all_items()\n",
    "    items = list(all_items)\n",
    "    if not items:\n",
    "        return None\n",
    "    \n",
    "    items = sorted(items, key=lambda x: x.properties.get('eo:cloud_cover', float('inf')))\n",
    "    lowest_cloud_item = items[0]\n",
    "\n",
    "    epsg = lowest_cloud_item.properties[\"proj:epsg\"]\n",
    "\n",
    "    poidf = gpd.GeoDataFrame(\n",
    "        pd.DataFrame(),\n",
    "        crs=\"EPSG:4326\",\n",
    "        geometry=[Point(lon, lat)],\n",
    "    ).to_crs(epsg)\n",
    "\n",
    "    coords = poidf.iloc[0].geometry.coords[0]\n",
    "\n",
    "    size = 256\n",
    "    gsd = 10 if COLLECTION == \"sentinel-2-l2a\" else 30\n",
    "    bounds = (\n",
    "        coords[0] - (size * gsd) // 2,\n",
    "        coords[1] - (size * gsd) // 2,\n",
    "        coords[0] + (size * gsd) // 2,\n",
    "        coords[1] + (size * gsd) // 2,\n",
    "    )\n",
    "\n",
    "    stack = stackstac.stack(\n",
    "        lowest_cloud_item,\n",
    "        bounds=bounds,\n",
    "        snap_bounds=False,\n",
    "        epsg=epsg,\n",
    "        resolution=gsd,\n",
    "        dtype=\"float64\",\n",
    "        rescale=False,\n",
    "        fill_value=0.0,\n",
    "        assets=[\"blue\", \"green\", \"red\", \"nir\"],\n",
    "        resampling=Resampling.nearest,\n",
    "    )\n",
    "\n",
    "    stack = stack.compute()\n",
    "\n",
    "    items = []\n",
    "    dates = []\n",
    "    for item in all_items:\n",
    "        if item.datetime.date() not in dates:\n",
    "            items.append(item)\n",
    "            dates.append(item.datetime.date())\n",
    "\n",
    "    platform = \"sentinel-2-l2a\" if COLLECTION == \"sentinel-2-l2a\" else \"landsat-c2l1\"\n",
    "    mean = []\n",
    "    std = []\n",
    "    waves = []\n",
    "    for band in stack.band:\n",
    "        mean.append(metadata[platform].bands.mean[str(band.values)])\n",
    "        std.append(metadata[platform].bands.std[str(band.values)])\n",
    "        waves.append(metadata[platform].bands.wavelength[str(band.values)])\n",
    "\n",
    "    transform = v2.Compose([v2.Normalize(mean=mean, std=std)])\n",
    "\n",
    "    datetimes = stack.time.values.astype(\"datetime64[s]\").tolist()\n",
    "    times = [normalize_timestamp(dat) for dat in datetimes]\n",
    "    week_norm = [dat[0] for dat in times]\n",
    "    hour_norm = [dat[1] for dat in times]\n",
    "\n",
    "    latlons = [normalize_latlon(lat, lon)] * len(times)\n",
    "    lat_norm = [dat[0] for dat in latlons]\n",
    "    lon_norm = [dat[1] for dat in latlons]\n",
    "\n",
    "    pixels = torch.from_numpy(stack.data.astype(np.float32)).to(device)\n",
    "    pixels = transform(pixels)\n",
    "\n",
    "    batch_size = 16\n",
    "    num_batches = math.ceil(len(stack) / batch_size)\n",
    "    \n",
    "    embeddings_list = []\n",
    "    for i in range(num_batches):\n",
    "        start_idx = i * batch_size\n",
    "        end_idx = min((i + 1) * batch_size, len(stack))\n",
    "        \n",
    "        batch_pixels = pixels[start_idx:end_idx].to(device)\n",
    "        batch_time = torch.tensor(np.hstack((week_norm, hour_norm))[start_idx:end_idx], dtype=torch.float32).to(device)\n",
    "        batch_latlon = torch.tensor(np.hstack((lat_norm, lon_norm))[start_idx:end_idx], dtype=torch.float32).to(device)\n",
    "        \n",
    "        batch_datacube = {\n",
    "            \"platform\": platform,\n",
    "            \"time\": batch_time,\n",
    "            \"latlon\": batch_latlon,\n",
    "            \"pixels\": batch_pixels,\n",
    "            \"gsd\": torch.tensor(stack.gsd.values).to(device),\n",
    "            \"waves\": torch.tensor(waves).to(device),\n",
    "        }\n",
    "\n",
    "        batch_datacube = to_device(batch_datacube, device)\n",
    "\n",
    "        try:\n",
    "            model = model.to(device)\n",
    "\n",
    "            with torch.no_grad():\n",
    "                unmsk_patch, _, _, _ = model.model.encoder(batch_datacube)\n",
    "            batch_embeddings = unmsk_patch[:, 0, :].cpu().numpy()\n",
    "            embeddings_list.append(batch_embeddings)\n",
    "        except RuntimeError as e:\n",
    "            if \"out of memory\" in str(e):\n",
    "                print(f\"GPU OOM for point ({lon}, {lat}), batch {i+1}/{num_batches}. Trying CPU...\")\n",
    "                device = torch.device(\"cpu\")\n",
    "                batch_datacube = to_device(batch_datacube, device)\n",
    "                model = model.to(device)\n",
    "                with torch.no_grad():\n",
    "                    unmsk_patch, _, _, _ = model.model.encoder(batch_datacube)\n",
    "                batch_embeddings = unmsk_patch[:, 0, :].numpy()\n",
    "                embeddings_list.append(batch_embeddings)\n",
    "            else:\n",
    "                raise e\n",
    "\n",
    "    embeddings = np.concatenate(embeddings_list, axis=0)\n",
    "    return embeddings\n",
    "\n",
    "def create_embeddings(start_date, end_date, grid, COLLECTION, start_date2=None, end_date2=None):\n",
    "    torch.cuda.empty_cache()  # Clear GPU cache\n",
    "    \n",
    "    # Assuming grid is a GeoDataFrame with the points, get a list of points in EPSG:4326\n",
    "    points = grid.to_crs(\"EPSG:4326\").geometry.apply(lambda x: (x.x, x.y)).tolist()\n",
    "\n",
    "    # Store results in a list\n",
    "    results = []\n",
    "\n",
    "    # Iterate through the points and process each one\n",
    "    for i, point in enumerate(tqdm(points)):\n",
    "        lon, lat = point\n",
    "        embeddings = process_point(lon, lat, model, metadata, start_date, end_date, COLLECTION, device=\"cuda\")\n",
    "        \n",
    "        # Initialize embeddings_new as None\n",
    "        embeddings_new = None\n",
    "        if start_date2 is not None and end_date2 is not None:\n",
    "            embeddings_new = process_point(lon, lat, model, metadata, start_date2, end_date2, COLLECTION, device=\"cuda\")\n",
    "        \n",
    "        if embeddings is not None:\n",
    "            if embeddings_new is not None:\n",
    "                results.append((lon, lat, embeddings, embeddings_new))\n",
    "            else:\n",
    "                results.append((lon, lat, embeddings))\n",
    "\n",
    "    # Define the columns based on whether embeddings_new is used\n",
    "    columns = [\"lon\", \"lat\", \"embeddings\"]\n",
    "    if start_date2 is not None and end_date2 is not None:\n",
    "        columns.append(\"embeddings_new\")\n",
    "\n",
    "    # Create a DataFrame from the results\n",
    "    df = pd.DataFrame(results, columns=columns)\n",
    "\n",
    "    # Convert to a GeoDataFrame\n",
    "    gdf_results = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.lon, df.lat))\n",
    "\n",
    "    return gdf_results\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd106ff-00c3-41cd-832c-f30c4734bbfb",
   "metadata": {},
   "source": [
    "#### Utility Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62771797-ccdc-4a07-8bce-894aceb802a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility functions for example models\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from sklearn.metrics import mean_squared_error, f1_score, accuracy_score, root_mean_squared_error, mean_absolute_error\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Dense, Dropout, Input\n",
    "from keras.optimizers import Adam\n",
    "import pyproj\n",
    "from datetime import datetime\n",
    "import rioxarray\n",
    "\n",
    "# Detect if GPU is available\n",
    "device = '/GPU:0' if tf.config.list_physical_devices('GPU') else '/CPU:0'\n",
    "\n",
    "class NeuralNetworkModel:\n",
    "    def __init__(self, model_type, input_shape, layers, dropout_rate, learning_rate, num_classes=None, device='/cpu:0'):\n",
    "        self.model_type = model_type\n",
    "        self.input_shape = input_shape\n",
    "        self.layers = layers\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.learning_rate = learning_rate\n",
    "        self.num_classes = num_classes\n",
    "        self.device = device\n",
    "        self.model = self.build_model()\n",
    "\n",
    "    def build_model(self):\n",
    "        model = Sequential()\n",
    "        model.add(Input(shape=(self.input_shape,)))\n",
    "        for layer_size in self.layers:\n",
    "            model.add(Dense(layer_size, activation='relu'))\n",
    "            model.add(Dropout(self.dropout_rate))\n",
    "\n",
    "        if self.model_type == 'regression':\n",
    "            model.add(Dense(1, activation='linear'))\n",
    "        elif self.model_type == 'multi_regression':\n",
    "            model.add(Dense(2, activation='linear'))\n",
    "        elif self.model_type == 'binary_classifier':\n",
    "            model.add(Dense(1, activation='sigmoid'))\n",
    "        elif self.model_type == 'multi_classifier':\n",
    "            model.add(Dense(self.num_classes, activation='softmax'))\n",
    "        \n",
    "        return model\n",
    "\n",
    "    def compile_model(self):\n",
    "        optimizer = Adam(learning_rate=self.learning_rate)\n",
    "        if self.model_type == 'regression':\n",
    "            self.model.compile(optimizer=optimizer, loss='mean_squared_error', metrics=['mae'])\n",
    "        elif self.model_type == 'binary_classifier':\n",
    "            self.model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=['accuracy'])\n",
    "        elif self.model_type == 'multi_classifier':\n",
    "            self.model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy'])\n",
    "\n",
    "    def train_model(self, X_train, y_train, epochs=20, batch_size=32, validation_split=0.2):\n",
    "        with tf.device(self.device):\n",
    "            history = self.model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, validation_split=validation_split)\n",
    "        return history\n",
    "\n",
    "    def evaluate_model(self, X_test, y_test):\n",
    "        with tf.device(self.device):\n",
    "            results = self.model.evaluate(X_test, y_test)\n",
    "        return results\n",
    "\n",
    "    def predict(self, X_test):\n",
    "        with tf.device(self.device):\n",
    "            predictions = self.model.predict(X_test)\n",
    "        \n",
    "        if self.model_type == 'binary_classifier':\n",
    "            return (predictions > 0.5).astype(int).flatten()\n",
    "        elif self.model_type == 'multi_classifier':\n",
    "            return np.argmax(predictions, axis=1)\n",
    "        return predictions\n",
    "\n",
    "    def calculate_metrics(self, y_test, predictions):\n",
    "        if self.model_type == 'regression':\n",
    "            return np.sqrt(mean_squared_error(y_test, predictions, multioutput='raw_values'))\n",
    "        elif self.model_type == 'binary_classifier':\n",
    "            return accuracy_score(y_test, predictions)\n",
    "        elif self.model_type == 'multi_classifier':\n",
    "            return f1_score(np.argmax(y_test, axis=1), predictions, average='weighted')\n",
    "    \n",
    "    def save_model(self, filename):\n",
    "        self.model.save(filename)\n",
    "\n",
    "    @classmethod\n",
    "    def load_model(cls, filename, model_type, input_shape, num_classes=None, device='/cpu:0'):\n",
    "        loaded_model = tf.keras.models.load_model(filename)\n",
    "        nn = cls(model_type, input_shape, [], 0, 0, num_classes, device)  # Dummy values for layers, dropout_rate, and learning_rate\n",
    "        nn.model = loaded_model\n",
    "        return nn\n",
    "\n",
    "def get_utm_zone(longitude):\n",
    "    return int((longitude + 180) / 6) + 1\n",
    "\n",
    "def gdf_preprocess(filepath, COLLECTION):\n",
    "    embedding_width = 2560 if COLLECTION == \"sentinel-2-l2a\" else 7680\n",
    "    gdf = gpd.read_file(filepath)\n",
    "\n",
    "    # Get the bounds of the geometry\n",
    "    minx, miny, maxx, maxy = gdf.geometry.bounds.iloc[0]\n",
    "\n",
    "    # Calculate UTM zone\n",
    "    utm_zone = get_utm_zone(minx)\n",
    "\n",
    "    # Get the corresponding EPSG code for the UTM zone using pyproj\n",
    "    utm_crs = pyproj.CRS(f\"+proj=utm +zone={utm_zone} +datum=WGS84\")\n",
    "    epsg_code = utm_crs.to_epsg()\n",
    "\n",
    "    # Reproject the GeoDataFrame to the chosen EPSG code\n",
    "    gdf = gdf.to_crs(epsg=epsg_code)\n",
    "\n",
    "    # Create a grid of points embedding_width meters apart\n",
    "    x = np.arange(gdf.total_bounds[0], gdf.total_bounds[2], embedding_width)\n",
    "    y = np.arange(gdf.total_bounds[1], gdf.total_bounds[3], embedding_width)\n",
    "    xx, yy = np.meshgrid(x, y)\n",
    "    points = np.vstack([xx.ravel(), yy.ravel()]).T\n",
    "\n",
    "    # Create GeoDataFrame from points\n",
    "    grid = gpd.GeoDataFrame(geometry=gpd.points_from_xy(points[:, 0], points[:, 1], crs=gdf.crs))\n",
    "    \n",
    "    # Calculate bounding boxes around each point\n",
    "    grid['bounding_box'] = grid.geometry.apply(lambda point: box(\n",
    "        point.x - embedding_width / 2,\n",
    "        point.y - embedding_width / 2,\n",
    "        point.x + embedding_width / 2,\n",
    "        point.y + embedding_width / 2\n",
    "    ))\n",
    "\n",
    "    return grid\n",
    "\n",
    "def final_grid(filepath, gsd, gdf_data):\n",
    "    gdf = gpd.read_file(filepath)\n",
    "\n",
    "    # Get the bounds of the geometry\n",
    "    minx, miny, maxx, maxy = gdf.geometry.bounds.iloc[0]\n",
    "\n",
    "    # Calculate UTM zone\n",
    "    utm_zone = get_utm_zone(minx)\n",
    "\n",
    "    # Get the corresponding EPSG code for the UTM zone using pyproj\n",
    "    utm_crs = pyproj.CRS(f\"+proj=utm +zone={utm_zone} +datum=WGS84\")\n",
    "    epsg_code = utm_crs.to_epsg()\n",
    "\n",
    "    # Reproject the GeoDataFrame to the chosen EPSG code\n",
    "    gdf = gdf.to_crs(epsg=epsg_code)\n",
    "\n",
    "    # Create a grid of points 5120m apart\n",
    "    x = np.arange(gdf.total_bounds[0], gdf.total_bounds[2], gsd)\n",
    "    y = np.arange(gdf.total_bounds[1], gdf.total_bounds[3], gsd)\n",
    "    xx, yy = np.meshgrid(x, y)\n",
    "    points = np.vstack([xx.ravel(), yy.ravel()]).T\n",
    "\n",
    "    grid = gpd.GeoDataFrame(geometry=gpd.points_from_xy(points[:, 0], points[:, 1], crs=gdf.crs))\n",
    "\n",
    "    # Ensure both GeoDataFrames have the same CRS\n",
    "    grid = grid.to_crs(4326)\n",
    "    gdf_data = gdf_data.to_crs(4326)\n",
    "\n",
    "    # Perform a spatial join to assign 'pred' values from gdf to gdf_final based on geometry intersection\n",
    "    gdf_final = gpd.sjoin(\n",
    "        grid, \n",
    "        gdf_data[['geometry'] + gdf_data.filter(like='pred').columns.tolist()], \n",
    "        how='left', \n",
    "        op='intersects'\n",
    "    )\n",
    "\n",
    "    # Drop the columns that were added during the join and rename the 'pred' column if necessary\n",
    "    gdf_final = gdf_final.drop(columns=['index_right'])\n",
    "\n",
    "    return(gdf_final)\n",
    "\n",
    "def NN_preds(task_id, model_type, gdf):\n",
    "    # Download model (as pickle file or similar)\n",
    "    loaded_nn = NeuralNetworkModel.load_model(f'models/task_{task_id}_model.h5', model_type=model_type, input_shape=768, num_classes=11, device=device)\n",
    "\n",
    "    # perform inference\n",
    "    new_data = np.squeeze(gdf['embeddings'].tolist())\n",
    "    new_data = pd.DataFrame(new_data)  # Ensure the new data is in DataFrame format\n",
    "\n",
    "    # Standardize the new data using the saved scaler\n",
    "    if os.path.exists(f'models/{task_id}_scaler.joblib'):\n",
    "        scaler = joblib.load(f'models/{task_id}_scaler.joblib')\n",
    "        new_data = scaler.transform(new_data)\n",
    "\n",
    "    # Make predictions\n",
    "    new_predictions = loaded_nn.predict(new_data)\n",
    "\n",
    "    # Load the label encoder\n",
    "    if os.path.exists(f'models/{task_id}_label_encoder.joblib'):\n",
    "        label_encoder = joblib.load(f'models/{task_id}_label_encoder.joblib')\n",
    "        gdf['pred'] = label_encoder.inverse_transform(new_predictions)\n",
    "    else:\n",
    "        if model_type == 'multi_regression':\n",
    "            gdf[['pred_maize','pred_wheat']] = new_predictions\n",
    "        else:\n",
    "            gdf['pred'] = new_predictions\n",
    "    return(gdf)\n",
    "\n",
    "def prep_gdf_for_sim_search(gdf, geojson_of_examples):\n",
    "    # For polygons in the gdf, a positive label is assigned if it contains a point from the geojson_of_examples\n",
    "    # OR at least 40% of the polygon is covered by a polygon from the geojson_of_examples\n",
    "\n",
    "    # Load the GeoDataFrame from the geojson file\n",
    "    examples_gdf = gpd.read_file(geojson_of_examples)\n",
    "\n",
    "    # Assign a positive label to the examples\n",
    "    examples_gdf['label'] = 'positive'\n",
    "\n",
    "    # Assign a negative label to the gdf\n",
    "    gdf['label'] = 'negative'\n",
    "\n",
    "    # Iterate through the examples and assign positive labels to the gdf\n",
    "    for i, example in examples_gdf.iterrows():\n",
    "        # Check if the example is a point or a polygon\n",
    "        if example['geometry'].type == 'Point':\n",
    "            # Assign a positive label to the point if it is contained within the gdf\n",
    "            gdf.loc[gdf.contains(example['geometry']), 'label'] = 'positive'\n",
    "        elif example['geometry'].type == 'Polygon':\n",
    "            # Calculate the intersection area\n",
    "            intersection_area = gdf.intersection(example['geometry']).area\n",
    "            # Calculate the coverage ratio\n",
    "            coverage_ratio = intersection_area / gdf.area\n",
    "            # Assign a positive label to geometries where the example polygon covers at least 40%\n",
    "            gdf.loc[(coverage_ratio >= 0.4) & (intersection_area > 0), 'label'] = 'positive'\n",
    "\n",
    "    return gdf\n",
    "\n",
    "def calculate_sim_score(example_embedding, positive_embeddings, negative_embeddings):\n",
    "    # Calculate the similarity score between the embedding and the positive and negative embeddings\n",
    "    # Uses the sum of the cosine similarities\n",
    "    sim_score = 0\n",
    "\n",
    "    for pos_embedding in positive_embeddings:\n",
    "        sim_score += np.dot(example_embedding, pos_embedding) / (np.linalg.norm(example_embedding) * np.linalg.norm(pos_embedding))\n",
    "\n",
    "    for neg_embedding in negative_embeddings:\n",
    "        sim_score -= np.dot(example_embedding, neg_embedding) / (np.linalg.norm(example_embedding) * np.linalg.norm(neg_embedding))\n",
    "\n",
    "    return sim_score \n",
    "\n",
    "def simulate_human_in_loop_search(gdf):\n",
    "    # Given a gdf with labels and embeddings, simulate a human-in-the-loop search\n",
    "    true_positive_examples = gdf[gdf['label'] == 'positive']\n",
    "    true_negative_examples = gdf[gdf['label'] == 'negative']\n",
    "\n",
    "    chosen_positive_embeddings = true_positive_examples.sample(n=10)[\"embeddings\"].tolist()\n",
    "    chosen_negative_embeddings = true_negative_examples.sample(n=10)[\"embeddings\"].tolist()\n",
    "\n",
    "    current_recall = len(chosen_positive_embeddings) / len(true_positive_examples)\n",
    "    rounds = 0\n",
    "\n",
    "    while current_recall < 0.9:  # Label the top 10 examples until recall is 90%\n",
    "        # Calculate the similarity score for each example\n",
    "        gdf['sim_score'] = gdf['embeddings'].apply(lambda x: calculate_sim_score(x, chosen_positive_embeddings, chosen_negative_embeddings))\n",
    "\n",
    "        # Label the top 10 examples, and exclude the already chosen examples\n",
    "        top_10 = gdf.sort_values('sim_score', ascending=False)\n",
    "        top_10 = top_10[~top_10['embeddings'].isin(chosen_positive_embeddings + chosen_negative_embeddings)].head(10)\n",
    "\n",
    "        # Update the chosen positive embeddings\n",
    "        chosen_positive_embeddings = chosen_positive_embeddings + top_10[top_10['label'] == 'positive'][\"embeddings\"].tolist()\n",
    "        chosen_negative_embeddings = chosen_negative_embeddings + top_10[top_10['label'] == 'negative'][\"embeddings\"].tolist()\n",
    "\n",
    "        # Update the current recall\n",
    "        current_recall = len(chosen_positive_embeddings) / len(true_positive_examples)\n",
    "        rounds += 1\n",
    "    \n",
    "    return rounds\n",
    "\n",
    "def process_geometries(combined_gdf, callable):\n",
    "    # Get the total number of rows in the GeoDataFrame\n",
    "    total_rows = len(combined_gdf)\n",
    "\n",
    "    # Determine the number of threads to use\n",
    "    max_threads = 10  # Adjust this based on your system and Earth Engine quota\n",
    "    num_threads = min(total_rows, max_threads)\n",
    "\n",
    "    results = []\n",
    "\n",
    "    with ThreadPoolExecutor(max_workers=num_threads) as executor:\n",
    "        # Submit all tasks\n",
    "        future_to_index = {executor.submit(callable, row.geometry): index \n",
    "                        for index, row in combined_gdf.iterrows()}\n",
    "        \n",
    "        # Process as they complete with a progress bar\n",
    "        with tqdm(total=total_rows, desc=\"Processing bounding boxes\") as pbar:\n",
    "            for future in as_completed(future_to_index):\n",
    "                index = future_to_index[future]\n",
    "                try:\n",
    "                    result = future.result()\n",
    "                except Exception as exc:\n",
    "                    print(f'Generated an exception: {exc}')\n",
    "                    result = None\n",
    "                results.append((index, result))\n",
    "                pbar.update(1)\n",
    "\n",
    "    # Sort results by index and extract only the values\n",
    "    sorted_results = [r[1] for r in sorted(results, key=lambda x: x[0])]\n",
    "    \n",
    "    return sorted_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18ebb80f-fc81-48dc-ab2e-b6775a2b621f",
   "metadata": {},
   "source": [
    "#### Submission Information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bad12e6-e879-47a4-974a-1e1db39624ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "submission = {}\n",
    "\n",
    "# Participant data to fill\n",
    "submission['eligibility'] = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a253cb4-c74f-4fa2-98e1-9e1028c6631a",
   "metadata": {},
   "source": [
    "#### Task: 1 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df1afdf2-28f3-4e2b-a230-97a6907d141a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task 1 Specific Utilities\n",
    "\n",
    "import ee\n",
    "from tqdm import tqdm\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "ee.Initialize()\n",
    "lulc = ee.ImageCollection(\"projects/sat-io/open-datasets/landcover/ESRI_Global-LULC_10m_TS\").mosaic()\n",
    "\n",
    "def get_most_common_lulc(point):\n",
    "    # Convert the GeoPandas geometry to an Earth Engine geometry\n",
    "    ee_geometry = ee.Geometry.Point(point.x, point.y)\n",
    "    \n",
    "    # Get the LULC values within the bounding box\n",
    "    lulc_values = lulc.reduceRegion(\n",
    "        reducer=ee.Reducer.frequencyHistogram(),\n",
    "        geometry=ee_geometry,\n",
    "        scale=10,\n",
    "        maxPixels=1e9\n",
    "    ).get('b1')\n",
    "    \n",
    "    # Find the most common LULC value\n",
    "    lulc_dict = ee.Dictionary(lulc_values)\n",
    "    most_common = lulc_dict.keys().sort(lulc_dict.values()).get(-1)\n",
    "    \n",
    "    # Return the result\n",
    "    return int(most_common.getInfo())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35a60644-ecd4-47f9-8fdc-2bb13abc16ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION=\"sentinel-2-l2a\"\n",
    "\n",
    "# Load the GeoJSON file and generate grid of image chips\n",
    "geojson_path = 'test_data/challenge_1_bb.geojson'\n",
    "grid = gdf_preprocess(geojson_path, COLLECTION)\n",
    "\n",
    "# Create embeddings over the defined grid\n",
    "gdf = create_embeddings(start_date=\"2022-01-01\", end_date=\"2022-12-31\", grid=grid, COLLECTION=COLLECTION)\n",
    "\n",
    "###\n",
    "### replace w/ submission model code ###\n",
    "\n",
    "gdf = NN_preds(task_id=1, model_type=\"multi_classifier\", gdf=gdf)\n",
    "\n",
    "### replace w/ submission model code ###\n",
    "###\n",
    "\n",
    "# Load testing values from raster\n",
    "gdf['geometry'] = grid.geometry\n",
    "test_values = process_geometries(gdf.to_crs(4236), get_most_common_lulc)\n",
    "\n",
    "# calculate score\n",
    "submission['task_1_score'] = f1_score(test_values, gdf[\"pred\"], average=\"micro\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bd56b64-e07a-45e1-a566-903f5b472187",
   "metadata": {},
   "source": [
    "#### Task: 2 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d365a2e-4e3f-4ceb-a19a-6e54d515802e",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION=\"sentinel-2-l2a\"\n",
    "\n",
    "aquaculture_geojson_ns = 'test_data/challenge_2_ns.geojson'\n",
    "# aquaculture_geojson_us = 'test_data/challenge_2_us.geojson'\n",
    "# aquaculture_geojson_norway = 'test_data/challenge_2_norway.geojson'\n",
    "\n",
    "aquaculture_geojsons = [aquaculture_geojson_ns]\n",
    "\n",
    "gdfs = []\n",
    "\n",
    "for geojson in aquaculture_geojsons:\n",
    "\n",
    "    # Load the GeoJSON file\n",
    "    geojson_path = geojson\n",
    "    grid = gdf_preprocess(geojson_path, COLLECTION)\n",
    "\n",
    "    # Create embeddings over the defined grid\n",
    "    gdf = create_embeddings(start_date=\"2022-01-01\", end_date=\"2022-12-31\", grid=grid, COLLECTION=COLLECTION)\n",
    "\n",
    "    # Adds positive and negative labels to the gdf according to the presence of Aquaculture sites\n",
    "    gdf = prep_gdf_for_sim_search(gdf, geojson)\n",
    "\n",
    "    gdfs.append(gdf)\n",
    "\n",
    "\n",
    "gdf = pd.concat(gdfs)\n",
    "\n",
    "results = []\n",
    "for i in range(100):\n",
    "    rounds = simulate_human_in_loop_search(gdf)\n",
    "    results.append(rounds)\n",
    "\n",
    "\n",
    "task_2_score = sum(results) / len(results)\n",
    "submission['task_2_score'] = task_2_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fc4122e-96e1-447a-9013-eea0778ece68",
   "metadata": {},
   "source": [
    "#### Task: 3 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57446502-1505-4539-9036-398923438064",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task 3 Specific Utilities\n",
    "\n",
    "# Load the NASA/ORNL biomass carbon density dataset\n",
    "biomass = ee.ImageCollection(\"NASA/ORNL/biomass_carbon_density/v1\").mosaic()\n",
    "\n",
    "def get_agb(geometry):\n",
    "    ee_geometry = ee.Geometry.Polygon(list(geometry.exterior.coords))\n",
    "\n",
    "    agb_value = biomass.reduceRegion(\n",
    "        reducer=ee.Reducer.mean(),\n",
    "        geometry=ee_geometry,\n",
    "        scale=300,\n",
    "        maxPixels=1e9\n",
    "    ).get('agb')\n",
    "    \n",
    "    # Return the result\n",
    "    return agb_value.getInfo() if agb_value is not None else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbabd846-90f4-4af6-8a8d-77bd8b60b9ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION=\"sentinel-2-l2a\"\n",
    "\n",
    "# Load the GeoJSON file and generate grid of image chips\n",
    "geojson_path = 'test_data/challenge_3_bb.geojson'\n",
    "grid = gdf_preprocess(geojson_path, COLLECTION=COLLECTION)\n",
    "\n",
    "# Create embeddings over the defined grid\n",
    "gdf = create_embeddings(start_date=\"2022-01-01\", end_date=\"2022-12-31\", grid=grid, COLLECTION=COLLECTION)\n",
    "\n",
    "###\n",
    "### replace w/ submission model code ### ###\n",
    "\n",
    "gdf = NN_preds(task_id=3, model_type=\"regression\", gdf=gdf)\n",
    "\n",
    "### replace w/ submission model code ### ###\n",
    "###\n",
    "\n",
    "# Load testing values\n",
    "gdf['geometry'] = grid[\"bounding_box\"]\n",
    "test_values = process_geometries(gdf.to_crs(4236), get_agb)\n",
    "\n",
    "# calculate score\n",
    "submission['task_3_score'] = root_mean_squared_error(test_values, gdf[\"pred\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "606c6b00-8730-4645-b9a3-27e62d9caeab",
   "metadata": {},
   "source": [
    "#### Task: 4 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b422d10-c71f-4c1d-8df4-061c3bc156ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task 4 Specific Utilities\n",
    "\n",
    "TEST_START_DATE = os.getenv(\"TEST_START_DATE\")\n",
    "TEST_END_DATE = os.getenv(\"TEST_END_DATE\")\n",
    "TEST_YEAR = datetime.strptime(TEST_START_DATE, \"%Y-%m-%d\").year\n",
    "\n",
    "FLOOD_FILE = f\"./test_data/MaizeCWP/MaizeCWP{TEST_YEAR}.tif\"\n",
    "MAIZE_RASTER = rioxarray.open_rasterio(FLOOD_FILE).rio.reproject(4326)\n",
    "\n",
    "def get_crop_yield_at_point(geometry):\n",
    "    try:\n",
    "        clipped = MAIZE_RASTER.rio.clip([geometry], MAIZE_RASTER.rio.crs)\n",
    "        return float(clipped.sum().values)\n",
    "    except rioxarray.exceptions.NoDataInBounds as e:\n",
    "        return np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b1ca520-dcc2-4a10-84a7-f6408d568ecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION=\"landsat-c2-l2\"\n",
    "\n",
    "# Load the GeoJSON file and generate grid of image chips\n",
    "geojson_path = 'test_data/challenge_4_bb.geojson'\n",
    "grid = gdf_preprocess(geojson_path, COLLECTION=COLLECTION)\n",
    "\n",
    "# Create embeddings over the defined grid\n",
    "gdf = create_embeddings(start_date=TEST_START_DATE, end_date=TEST_END_DATE, grid=grid, COLLECTION=COLLECTION)\n",
    "\n",
    "###\n",
    "### replace w/ submission model code ### ###\n",
    "\n",
    "gdf = NN_preds(task_id=4, model_type=\"multi_regression\", gdf=gdf)\n",
    "\n",
    "### replace w/ submission model code ### ###\n",
    "###\n",
    "\n",
    "# Load testing values\n",
    "gdf['geometry'] = grid['bounding_box']\n",
    "test_values = [0 if i == -np.inf else i for i in process_geometries(gdf.to_crs(4236), get_crop_yield_at_point)]\n",
    "\n",
    "# calculate score\n",
    "submission['task_4_score'] = root_mean_squared_error(test_values, gdf[\"pred\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b36a66e6-9896-4384-9703-ca124562c04a",
   "metadata": {},
   "source": [
    "#### Task: 5 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1723ef7-a43b-46f4-bb00-07d370561f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task 5 Specific Utilities\n",
    "\n",
    "# Load the TIFF files as xarray Datasets\n",
    "FLOOD_PATH = os.getenv(\"FLOOD_PATH\")\n",
    "FLOOD_RASTER = rioxarray.open_rasterio(FLOOD_PATH)\n",
    "\n",
    "def get_flood_value_at_point(geometry):\n",
    "    try:\n",
    "        clipped = FLOOD_RASTER.rio.clip([geometry], FLOOD_RASTER.rio.crs)\n",
    "        return float(clipped.mean().values)\n",
    "    except rioxarray.exceptions.NoDataInBounds as e:\n",
    "        return np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aa800c4-5ef4-4a53-8557-63ce7b156b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION=\"sentinel-2-l2a\"\n",
    "\n",
    "# Load the GeoJSON file and generate grid of image chips\n",
    "geojson_path = 'test_data/challenge_5_bb.geojson'\n",
    "grid = gdf_preprocess(geojson_path, COLLECTION=COLLECTION)\n",
    "\n",
    "# Create embeddings over the defined grid\n",
    "gdf = create_embeddings(start_date=\"2018-08-01\", end_date=\"2018-09-30\", start_date2=\"2019-08-01\", end_date2=\"2019-09-30\", grid=grid, COLLECTION=COLLECTION) \n",
    "gdf['embeddings_delta'] = gdf['embeddings_new'] - gdf['embeddings']\n",
    "\n",
    "###\n",
    "### replace w/ submission model code ### ###\n",
    "\n",
    "gdf = NN_preds(task_id=5, model_type=\"binary_classifier\", gdf=gdf)\n",
    "\n",
    "### replace w/ submission model code ### ###\n",
    "###\n",
    "\n",
    "# Load testing values\n",
    "gdf['geometry'] = grid.geometry\n",
    "test_values = process_geometries(gdf_final, get_flood_value_at_point)\n",
    "\n",
    "# calculate score\n",
    "submission['task_5_score'] = f1_score(test_values, gdf_final[\"pred\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0fe75df-c323-4a60-aca4-e0f6e6c9d1d1",
   "metadata": {},
   "source": [
    "#### Task: 6 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec4099f6-816d-4c02-be5a-b4c4e20d8b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task 6 Specific Utilities\n",
    "\n",
    "def get_severity(geometry):\n",
    "    ee_geometry = ee.Geometry.Polygon(list(geometry.exterior.coords))\n",
    "\n",
    "    severity_value = burn_severity.reduceRegion(\n",
    "        reducer=ee.Reducer.mean(),\n",
    "        geometry=ee_geometry,\n",
    "        scale=30,\n",
    "        maxPixels=1e9\n",
    "    ).get('Severity') \n",
    "    \n",
    "    return severity_value.getInfo() if severity_value is not None else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "410fe96e-0d89-455d-a6bf-123ac9f665b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLLECTION=\"sentinel-2-l2a\"\n",
    "\n",
    "# Load the GeoJSON file and generate grid of image chips\n",
    "geojson_path = 'test_data/challenge_6_bb.geojson'\n",
    "grid = gdf_preprocess(geojson_path, COLLECTION=COLLECTION)\n",
    "\n",
    "# Create embeddings over the defined grid\n",
    "gdf = create_embeddings(start_date=\"2019-01-01\", end_date=\"2019-12-31\", start_date2=\"2020-01-01\", end_date2=\"2020-12-31\", grid=grid, COLLECTION=COLLECTION) \n",
    "gdf['embeddings_delta'] = gdf['embeddings_new'] - gdf['embeddings']\n",
    "\n",
    "###\n",
    "### replace w/ submission model code ### ###\n",
    "\n",
    "gdf = NN_preds(task_id=6, model_type=\"regression\", gdf=gdf)\n",
    "\n",
    "### replace w/ submission model code ### ###\n",
    "###\n",
    "\n",
    "# Load testing values\n",
    "gdf['geometry'] = grid['bounding_box']\n",
    "test_values = process_geometries(gdf.to_crs(4236), get_severity)\n",
    "\n",
    "# calculate score\n",
    "submission['task_6_score'] = root_mean_squared_error(test_values, gdf_final[\"pred\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "798b56fa-244f-4d9a-88f8-744acd7e9751",
   "metadata": {},
   "source": [
    "#### Task: 7 Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2305bcf4-1d1c-411a-a6c5-cc458b072fa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is a good example, but is currently broken. The tensors need more information to be processed correctly, and\n",
    "# this fix will be coming very soon.\n",
    "\n",
    "class CloudMaskingDataset(Dataset):\n",
    "    def __init__(self, chip_catalog, cloud_mask_dir, hls_chip_dir, metadata, transform=None, train_or_validate=\"train\"):\n",
    "        self.chip_catalog = pd.read_csv(chip_catalog)\n",
    "        self.chip_catalog = self.chip_catalog[self.chip_catalog['usage'] == train_or_validate]\n",
    "\n",
    "        self.cloud_mask_dir = cloud_mask_dir\n",
    "\n",
    "        self.hls_chip_dir = hls_chip_dir\n",
    "        self.metadata = metadata\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.chip_catalog)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        chip_id = self.chip_catalog.iloc[idx]['chip_id']\n",
    "        chip_path = f\"{self.hls_chip_dir}/{chip_id}.tif\"\n",
    "        cloud_mask_paths = glob(f'{self.cloud_mask_dir}/{chip_id}_*.tif') + [\"None\"]\n",
    "\n",
    "        with rasterio.open(chip_path) as chip:\n",
    "            chip_data = chip.read()  # Shape: (18, 224, 224)\n",
    "\n",
    "        masked_chip_data = chip_data.copy()\n",
    "\n",
    "        for time_step in range(3):\n",
    "            chosen_mask_path = random.choice(cloud_mask_paths)\n",
    "            \n",
    "            if chosen_mask_path == \"None\":\n",
    "                cloud_mask = np.ones((224, 224), dtype=np.uint8)\n",
    "            else:\n",
    "                with rasterio.open(chosen_mask_path) as mask:\n",
    "                    cloud_mask = mask.read(1)\n",
    "\n",
    "            masked_chip_data[time_step*6:(time_step+1)*6, :, :] *= (cloud_mask == 0)\n",
    "\n",
    "        images = [masked_chip_data[i*6:(i+1)*6, :, :] for i in range(3)]\n",
    "        ground_truth = [chip_data[i*6:(i+1)*6, :, :] for i in range(3)]\n",
    "\n",
    "        if self.transform:\n",
    "            images = [self.transform(torch.from_numpy(img).float()) for img in images]\n",
    "            ground_truth = [torch.from_numpy(gt).float() for gt in ground_truth]\n",
    "\n",
    "        return torch.stack(images), torch.stack(ground_truth)\n",
    "\n",
    "class FineTuningModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(FineTuningModel, self).__init__()\n",
    "        self.fc = nn.Linear(768, 512 * 7 * 7)  # Fully connected layer to reshape the embedding\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose2d(512, 256, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.ConvTranspose2d(256, 128, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(128, 64, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.ConvTranspose2d(64, 32, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.ConvTranspose2d(32, 16, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.ConvTranspose2d(16, 6, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.ConvTranspose2d(6, 6, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.ConvTranspose2d(6, 18, 3, padding=1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc(x)  # Apply the fully connected layer\n",
    "        x = x.view(-1, 512, 7, 7)  # Reshape to prepare for transposed convolutions\n",
    "        x = self.decoder(x)  # Apply the decoder to upsample to the target size\n",
    "        return x\n",
    "\n",
    "metadata = Box(yaml.safe_load(open(\"configs/metadata.yaml\")))\n",
    "\n",
    "# Set up the normalization\n",
    "platform = \"sentinel-2-l2a\"  # Example platform\n",
    "sentinel_band_names = [\"blue\", \"green\", \"red\", \"nir\", \"swir16\", \"swir22\"]\n",
    "mean = [metadata[platform].bands.mean[sentinel_band_names[i]] for i in range(6)]\n",
    "std = [metadata[platform].bands.std[sentinel_band_names[i]] for i in range(6)]\n",
    "transform = v2.Compose([v2.Normalize(mean=mean, std=std)])\n",
    "\n",
    "# Create the dataset and dataloader\n",
    "dataset = CloudMaskingDataset(\n",
    "    chip_catalog=\"./train_data/challenge_7_data/chip_catalog.csv\",\n",
    "    cloud_mask_dir=\"./train_data/challenge_7_data/train/cloud_masks\",\n",
    "    hls_chip_dir=\"./train_data/challenge_7_data/train/hls_chips\",\n",
    "    metadata=metadata,\n",
    "    transform=transform,\n",
    "    train_or_validate=\"train\"\n",
    ")\n",
    "\n",
    "dataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "\n",
    "# Example training loop\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Load the embedding model. This will be YOUR foundation model.\n",
    "embedding_model = ClayMAEModule.load_from_checkpoint(  \n",
    "    ckpt, metadata_path=\"configs/metadata.yaml\", shuffle=False, mask_ratio=0\n",
    ")\n",
    "embedding_model.eval()\n",
    "embedding_model = embedding_model.to(device)\n",
    "\n",
    "fine_tuning_model = FineTuningModel().to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(fine_tuning_model.parameters(), lr=1e-3)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(50):\n",
    "    fine_tuning_model.train()\n",
    "    for masked_images, ground_truth_images in dataloader:\n",
    "        masked_images = masked_images.to(device)\n",
    "        ground_truth_images = ground_truth_images.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        outputs = []\n",
    "        for i in range(3):\n",
    "            with torch.no_grad():\n",
    "                embeddings = embedding_model.encoder(masked_images[:, i]).view(masked_images.size(0), -1)  # Flatten the embedding\n",
    "            reconstructed_image = fine_tuning_model(embeddings)  # Pass the embedding through the fine-tuning model\n",
    "            outputs.append(reconstructed_image)\n",
    "        \n",
    "        outputs = torch.stack(outputs, dim=1)\n",
    "        \n",
    "        loss = criterion(outputs, ground_truth_images)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        print(f\"Epoch {epoch + 1}, Loss: {loss.item()}\")\n",
    "\n",
    "\n",
    "# Validation Loop\n",
    "fine_tuning_model.eval()\n",
    "val_dataset = CloudMaskingDataset(\n",
    "    chip_catalog=\"./train_data/challenge_7_data/chip_catalog.csv\",\n",
    "    cloud_mask_dir=\"./train_data/challenge_7_data/train/cloud_masks\",\n",
    "    hls_chip_dir=\"./train_data/challenge_7_data/train/hls_chips\",\n",
    "    metadata=metadata,\n",
    "    transform=transform,\n",
    "    train_or_validate=\"validate\"\n",
    ")\n",
    "\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=16, shuffle=False)\n",
    "\n",
    "val_losses = []\n",
    "\n",
    "for masked_images, ground_truth_images in val_dataloader:\n",
    "    masked_images = masked_images.to(device)\n",
    "    ground_truth_images = ground_truth_images.to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = []\n",
    "        for i in range(3):\n",
    "            embeddings = embedding_model.encoder(masked_images[:, i]).view(masked_images.size(0), -1)\n",
    "            reconstructed_image = fine_tuning_model(embeddings)\n",
    "            outputs.append(reconstructed_image)\n",
    "\n",
    "        outputs = torch.stack(outputs, dim=1)\n",
    "\n",
    "        loss = criterion(outputs, ground_truth_images)\n",
    "        val_losses.append(loss.item())\n",
    "\n",
    "print(f\"Validation Loss: {np.mean(val_losses)}\")\n",
    "\n",
    "# Finally, export the results on the test set\n",
    "# This is a placeholder for the actual submission code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c759b81-a85d-446f-8a30-96c40a8e44fe",
   "metadata": {},
   "source": [
    "#### Output submissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a20a4e0e-eb99-4b81-a3d9-3ffdbe02110a",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame([submission]).to_csv('output.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}