downstream_task_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8f13dc11",
   "metadata": {},
   "source": [
    "# Using Presto for downstream tasks\n",
    "\n",
    "The purpose of this notebook is to demonstrate how Presto (and utility functions in the Presto package) can be used for downstream tasks.\n",
    "\n",
    "To demonsrate the usefulness of Presto even when the input looks very different then what Presto was pre-trained on, we will consider tree-type mapping using single-timestep images.\n",
    "\n",
    "To do this, we will use the [TreeSat](https://essd.copernicus.org/articles/15/681/2023/) benchmark dataset. This tutorial requires the S2 data to be downloaded from [Zenodo](https://zenodo.org/record/6780578) and unzipped in the [treesat folder](data/treesat)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b633f57b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/gabrieltseng/anaconda3/envs/lem/lib/python3.9/site-packages/geopandas/_compat.py:106: UserWarning: The Shapely GEOS version (3.11.1-CAPI-1.17.1) is incompatible with the GEOS version PyGEOS was compiled with (3.10.4-CAPI-1.16.2). Conversions between both will be slow.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import xarray\n",
    "from pyproj import Transformer\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import f1_score\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "import presto\n",
    "\n",
    "# this is to silence the xarray deprecation warning.\n",
    "# Our version of xarray is pinned, but we'll need to fix this\n",
    "# when we upgrade\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4246895f",
   "metadata": {},
   "source": [
    "If the TreeSat data has been correctly downloaded from Zenodo (see the Markdown cell above), these assert statements should pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a27dcd41",
   "metadata": {},
   "outputs": [],
   "source": [
    "treesat_folder = presto.utils.data_dir / \"treesat\"\n",
    "assert treesat_folder.exists()\n",
    "\n",
    "# this folder should exist once the s2 file from zenodo has been unzipped\n",
    "s2_data_60m = treesat_folder / \"s2/60m\"\n",
    "assert s2_data_60m.exists()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51d38221",
   "metadata": {},
   "source": [
    "For simplicity, we will only consider classification between 2 tree species: Abies_alba and Acer_pseudoplatanus.\n",
    "\n",
    "The TreeSatAI S2 data contains the following bands: [\"B2\", \"B3\", \"B4\", \"B8\", \"B5\", \"B6\", \"B7\", \"B8A\", \"B11\", \"B12\", \"B1\", \"B9\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "464ff6a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "TREESATAI_S2_BANDS = [\"B2\", \"B3\", \"B4\", \"B8\", \"B5\", \"B6\", \"B7\", \"B8A\", \"B11\", \"B12\", \"B1\", \"B9\"]\n",
    "\n",
    "SPECIES = [\"Abies_alba\", \"Acer_pseudoplatanus\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72d6575a",
   "metadata": {},
   "source": [
    "### 1. Processing the data\n",
    "\n",
    "The TreeSatAI data is stored in `.tif` files. We will extract 9 pixels (from the 36 total pixels in each tif file) to construct our input data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0667a6d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3430 train files and 375 test files\n"
     ]
    }
   ],
   "source": [
    "# takes a (6, 6) treesat tif file, and returns a\n",
    "# (9,1,18) cropharvest eo-style file (with all bands \"masked\"\n",
    "# except for S1 and S2)\n",
    "INDICES_IN_TIF_FILE = list(range(0, 6, 2))\n",
    "\n",
    "with (treesat_folder / \"train_filenames.lst\").open(\"r\") as f:\n",
    "    train_files = [line for line in f if (line.startswith(SPECIES[0]) or line.startswith(SPECIES[1]))]\n",
    "with (treesat_folder / \"test_filenames.lst\").open(\"r\") as f:\n",
    "    test_files = [line for line in f if (line.startswith(SPECIES[0]) or line.startswith(SPECIES[1]))]\n",
    "\n",
    "print(f\"{len(train_files)} train files and {len(test_files)} test files\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "551ec000",
   "metadata": {},
   "source": [
    "We use `presto.construct_single_presto_input` to transform the TreeSat S2 data into the tensors expected by Presto."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d2a7d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_images(filenames):\n",
    "    arrays, masks, latlons, image_names, labels, dynamic_worlds = [], [], [], [], [], []\n",
    "    \n",
    "    for filename in tqdm(filenames):\n",
    "        tif_file = xarray.open_rasterio(s2_data_60m / filename.strip())\n",
    "        crs = tif_file.crs.split(\"=\")[-1]\n",
    "        transformer = Transformer.from_crs(crs, \"EPSG:4326\", always_xy=True)\n",
    "        \n",
    "        for x_idx in INDICES_IN_TIF_FILE:\n",
    "            for y_idx in INDICES_IN_TIF_FILE:\n",
    "                \n",
    "                # firstly, get the latitudes and longitudes\n",
    "                x, y = tif_file.x[x_idx], tif_file.y[y_idx]\n",
    "                lon, lat = transformer.transform(x, y) \n",
    "                latlons.append(torch.tensor([lat, lon]))\n",
    "                \n",
    "                # then, get the eo_data, mask and dynamic world\n",
    "                s2_data_for_pixel = torch.from_numpy(tif_file.values[:, x_idx, y_idx].astype(int)).float()\n",
    "                s2_data_with_time_dimension = s2_data_for_pixel.unsqueeze(0)\n",
    "                x, mask, dynamic_world = presto.construct_single_presto_input(\n",
    "                    s2=s2_data_with_time_dimension, s2_bands=TREESATAI_S2_BANDS\n",
    "                )\n",
    "                arrays.append(x)\n",
    "                masks.append(mask)\n",
    "                dynamic_worlds.append(dynamic_world)\n",
    "                \n",
    "                labels.append(0 if filename.startswith(\"Abies\") else 1)\n",
    "                image_names.append(filename)\n",
    "\n",
    "    return (torch.stack(arrays, axis=0),\n",
    "            torch.stack(masks, axis=0),\n",
    "            torch.stack(dynamic_worlds, axis=0),\n",
    "            torch.stack(latlons, axis=0),\n",
    "            torch.tensor(labels),\n",
    "            image_names,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9440dbce",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3430/3430 [01:03<00:00, 54.26it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 375/375 [00:06<00:00, 62.25it/s]\n"
     ]
    }
   ],
   "source": [
    "train_data = process_images(train_files)\n",
    "test_data = process_images(test_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da4bed2c",
   "metadata": {},
   "source": [
    "### 2. Using Presto as a feature extractor for a random forest\n",
    "\n",
    "One way to use Presto is as a feature extractor for a simple model (e.g. a Random Forest). We do this below.\n",
    "\n",
    "We load the pretrained Presto model using `Presto.load_pretrained()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "47bfc821",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Presto(\n",
       "  (encoder): Encoder(\n",
       "    (eo_patch_embed): ModuleDict(\n",
       "      (S1): Linear(in_features=2, out_features=128, bias=True)\n",
       "      (S2_RGB): Linear(in_features=3, out_features=128, bias=True)\n",
       "      (S2_NIR_10m): Linear(in_features=1, out_features=128, bias=True)\n",
       "      (S2_NIR_20m): Linear(in_features=1, out_features=128, bias=True)\n",
       "      (S2_Red_Edge): Linear(in_features=3, out_features=128, bias=True)\n",
       "      (S2_SWIR): Linear(in_features=2, out_features=128, bias=True)\n",
       "      (ERA5): Linear(in_features=2, out_features=128, bias=True)\n",
       "      (SRTM): Linear(in_features=2, out_features=128, bias=True)\n",
       "      (NDVI): Linear(in_features=1, out_features=128, bias=True)\n",
       "    )\n",
       "    (dw_embed): Embedding(10, 128)\n",
       "    (latlon_embed): Linear(in_features=3, out_features=128, bias=True)\n",
       "    (blocks): ModuleList(\n",
       "      (0-1): 2 x Block(\n",
       "        (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=128, out_features=384, bias=True)\n",
       "          (q_norm): Identity()\n",
       "          (k_norm): Identity()\n",
       "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "          (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (ls1): Identity()\n",
       "        (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Mlp(\n",
       "          (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          (act): GELU(approximate='none')\n",
       "          (drop1): Dropout(p=0.0, inplace=False)\n",
       "          (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "          (drop2): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (ls2): Identity()\n",
       "      )\n",
       "    )\n",
       "    (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    (month_embed): Embedding(12, 32)\n",
       "    (channel_embed): Embedding(10, 32)\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (decoder_embed): Linear(in_features=128, out_features=128, bias=True)\n",
       "    (decoder_blocks): ModuleList(\n",
       "      (0-1): 2 x Block(\n",
       "        (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=128, out_features=384, bias=True)\n",
       "          (q_norm): Identity()\n",
       "          (k_norm): Identity()\n",
       "          (attn_drop): Dropout(p=0.0, inplace=False)\n",
       "          (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "          (proj_drop): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (ls1): Identity()\n",
       "        (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Mlp(\n",
       "          (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          (act): GELU(approximate='none')\n",
       "          (drop1): Dropout(p=0.0, inplace=False)\n",
       "          (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "          (drop2): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (ls2): Identity()\n",
       "      )\n",
       "    )\n",
       "    (decoder_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    (eo_decoder_pred): ModuleDict(\n",
       "      (S1): Linear(in_features=128, out_features=2, bias=True)\n",
       "      (S2_RGB): Linear(in_features=128, out_features=3, bias=True)\n",
       "      (S2_NIR_10m): Linear(in_features=128, out_features=1, bias=True)\n",
       "      (S2_NIR_20m): Linear(in_features=128, out_features=1, bias=True)\n",
       "      (S2_Red_Edge): Linear(in_features=128, out_features=3, bias=True)\n",
       "      (S2_SWIR): Linear(in_features=128, out_features=2, bias=True)\n",
       "      (ERA5): Linear(in_features=128, out_features=2, bias=True)\n",
       "      (SRTM): Linear(in_features=128, out_features=2, bias=True)\n",
       "      (NDVI): Linear(in_features=128, out_features=1, bias=True)\n",
       "    )\n",
       "    (dw_decoder_pred): Linear(in_features=128, out_features=9, bias=True)\n",
       "    (channel_embeddings): Embedding(10, 32)\n",
       "    (month_embed): Embedding(12, 48)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size = 64\n",
    "\n",
    "pretrained_model = presto.Presto.load_pretrained()\n",
    "pretrained_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82dfa681",
   "metadata": {},
   "source": [
    "We will start by constructing features for the training data, and using this to train a Random Forest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c06e50f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the treesat AI data was collected during the summer,\n",
    "# so we estimate the month to be 6 (July)\n",
    "month = torch.tensor([6] * train_data[0].shape[0]).long()\n",
    "\n",
    "dl = DataLoader(\n",
    "    TensorDataset(\n",
    "        train_data[0].float(),  # x\n",
    "        train_data[1].bool(),  # mask\n",
    "        train_data[2].long(),  # dynamic world\n",
    "        train_data[3].float(),  # latlons\n",
    "        month\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "506a1c05",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 483/483 [00:05<00:00, 92.62it/s]\n"
     ]
    }
   ],
   "source": [
    "features_list = []\n",
    "for (x, mask, dw, latlons, month) in tqdm(dl):\n",
    "    with torch.no_grad():\n",
    "        encodings = (\n",
    "            pretrained_model.encoder(\n",
    "                x, dynamic_world=dw, mask=mask, latlons=latlons, month=month\n",
    "            )\n",
    "            .cpu()\n",
    "            .numpy()\n",
    "        )\n",
    "        features_list.append(encodings)\n",
    "features_np = np.concatenate(features_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "870d977f",
   "metadata": {},
   "source": [
    "We use `features_np` to train a Random Forest classifier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fe116e42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>#sk-container-id-1 {color: black;background-color: white;}#sk-container-id-1 pre{padding: 0;}#sk-container-id-1 div.sk-toggleable {background-color: white;}#sk-container-id-1 label.sk-toggleable__label {cursor: pointer;display: block;width: 100%;margin-bottom: 0;padding: 0.3em;box-sizing: border-box;text-align: center;}#sk-container-id-1 label.sk-toggleable__label-arrow:before {content: \"▸\";float: left;margin-right: 0.25em;color: #696969;}#sk-container-id-1 label.sk-toggleable__label-arrow:hover:before {color: black;}#sk-container-id-1 div.sk-estimator:hover label.sk-toggleable__label-arrow:before {color: black;}#sk-container-id-1 div.sk-toggleable__content {max-height: 0;max-width: 0;overflow: hidden;text-align: left;background-color: #f0f8ff;}#sk-container-id-1 div.sk-toggleable__content pre {margin: 0.2em;color: black;border-radius: 0.25em;background-color: #f0f8ff;}#sk-container-id-1 input.sk-toggleable__control:checked~div.sk-toggleable__content {max-height: 200px;max-width: 100%;overflow: auto;}#sk-container-id-1 input.sk-toggleable__control:checked~label.sk-toggleable__label-arrow:before {content: \"▾\";}#sk-container-id-1 div.sk-estimator input.sk-toggleable__control:checked~label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 div.sk-label input.sk-toggleable__control:checked~label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 input.sk-hidden--visually {border: 0;clip: rect(1px 1px 1px 1px);clip: rect(1px, 1px, 1px, 1px);height: 1px;margin: -1px;overflow: hidden;padding: 0;position: absolute;width: 1px;}#sk-container-id-1 div.sk-estimator {font-family: monospace;background-color: #f0f8ff;border: 1px dotted black;border-radius: 0.25em;box-sizing: border-box;margin-bottom: 0.5em;}#sk-container-id-1 div.sk-estimator:hover {background-color: #d4ebff;}#sk-container-id-1 div.sk-parallel-item::after {content: \"\";width: 100%;border-bottom: 1px solid gray;flex-grow: 1;}#sk-container-id-1 div.sk-label:hover label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 div.sk-serial::before {content: \"\";position: absolute;border-left: 1px solid gray;box-sizing: border-box;top: 0;bottom: 0;left: 50%;z-index: 0;}#sk-container-id-1 div.sk-serial {display: flex;flex-direction: column;align-items: center;background-color: white;padding-right: 0.2em;padding-left: 0.2em;position: relative;}#sk-container-id-1 div.sk-item {position: relative;z-index: 1;}#sk-container-id-1 div.sk-parallel {display: flex;align-items: stretch;justify-content: center;background-color: white;position: relative;}#sk-container-id-1 div.sk-item::before, #sk-container-id-1 div.sk-parallel-item::before {content: \"\";position: absolute;border-left: 1px solid gray;box-sizing: border-box;top: 0;bottom: 0;left: 50%;z-index: -1;}#sk-container-id-1 div.sk-parallel-item {display: flex;flex-direction: column;z-index: 1;position: relative;background-color: white;}#sk-container-id-1 div.sk-parallel-item:first-child::after {align-self: flex-end;width: 50%;}#sk-container-id-1 div.sk-parallel-item:last-child::after {align-self: flex-start;width: 50%;}#sk-container-id-1 div.sk-parallel-item:only-child::after {width: 0;}#sk-container-id-1 div.sk-dashed-wrapped {border: 1px dashed gray;margin: 0 0.4em 0.5em 0.4em;box-sizing: border-box;padding-bottom: 0.4em;background-color: white;}#sk-container-id-1 div.sk-label label {font-family: monospace;font-weight: bold;display: inline-block;line-height: 1.2em;}#sk-container-id-1 div.sk-label-container {text-align: center;}#sk-container-id-1 div.sk-container {/* jupyter's `normalize.less` sets `[hidden] { display: none; }` but bootstrap.min.css set `[hidden] { display: none !important; }` so we also need the `!important` here to be able to override the default hidden behavior on the sphinx rendered scikit-learn.org. See: https://github.com/scikit-learn/scikit-learn/issues/21755 */display: inline-block !important;position: relative;}#sk-container-id-1 div.sk-text-repr-fallback {display: none;}</style><div id=\"sk-container-id-1\" class=\"sk-top-container\"><div class=\"sk-text-repr-fallback\"><pre>RandomForestClassifier(class_weight=&#x27;balanced&#x27;, random_state=42)</pre><b>In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook. <br />On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.</b></div><div class=\"sk-container\" hidden><div class=\"sk-item\"><div class=\"sk-estimator sk-toggleable\"><input class=\"sk-toggleable__control sk-hidden--visually\" id=\"sk-estimator-id-1\" type=\"checkbox\" checked><label for=\"sk-estimator-id-1\" class=\"sk-toggleable__label sk-toggleable__label-arrow\">RandomForestClassifier</label><div class=\"sk-toggleable__content\"><pre>RandomForestClassifier(class_weight=&#x27;balanced&#x27;, random_state=42)</pre></div></div></div></div></div>"
      ],
      "text/plain": [
       "RandomForestClassifier(class_weight='balanced', random_state=42)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = RandomForestClassifier(class_weight=\"balanced\", random_state=42)\n",
    "model.fit(features_np, train_data[4].numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e468d44",
   "metadata": {},
   "source": [
    "We can then use this trained random forest to make some predictions on the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "04a90f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the treesat AI data was collected during the summer,\n",
    "# so we estimate the month to be 6 (July)\n",
    "month = torch.tensor([6] * test_data[0].shape[0]).long()\n",
    "\n",
    "dl = DataLoader(\n",
    "    TensorDataset(\n",
    "        test_data[0].float(),  # x\n",
    "        test_data[1].bool(),  # mask\n",
    "        test_data[2].long(),  # dynamic world\n",
    "        test_data[3].float(),  # latlons\n",
    "        month\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "94521467",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 53/53 [00:01<00:00, 27.22it/s]\n"
     ]
    }
   ],
   "source": [
    "test_preds = []\n",
    "for (x, mask, dw, latlons, month) in tqdm(dl):\n",
    "    with torch.no_grad():\n",
    "        pretrained_model.eval()\n",
    "        encodings = (pretrained_model.encoder(\n",
    "            x, dynamic_world=dw, mask=mask, latlons=latlons, month=month)\n",
    "            .cpu()\n",
    "            .numpy()\n",
    "        )\n",
    "        test_preds.append(model.predict_proba(encodings))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46471233",
   "metadata": {},
   "source": [
    "We take the mode of the 9 pixels we processed for a single image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2510600e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pix_per_image = len(INDICES_IN_TIF_FILE) ** 2\n",
    "\n",
    "test_preds_np = np.concatenate(test_preds, axis=0)\n",
    "test_preds_np = np.reshape(\n",
    "    test_preds_np,\n",
    "    (int(len(test_preds_np) / pix_per_image), pix_per_image, test_preds_np.shape[-1]),\n",
    ")\n",
    "# then, take the mode of the model predictions\n",
    "test_preds_np_argmax = stats.mode(\n",
    "    np.argmax(test_preds_np, axis=-1), axis=1, keepdims=False\n",
    ")[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8117843b",
   "metadata": {},
   "source": [
    "And finally, we can compute the F1 score of the test predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "02b54c15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9681132732641782"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target = np.reshape(test_data[4], (int(len(test_data[4]) / pix_per_image), pix_per_image))[:, 0]\n",
    "\n",
    "f1_score(target, test_preds_np_argmax, average=\"weighted\")"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}