{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Setup\n",
    "\n",
    "In colab:\n",
    "1. Go to \"Runtime\" -> \"Change runtime type\" -> Select \"T4 GPU\"\n",
    "2. Install TerraTorch"
   ],
   "id": "be93ad48ee7c0f29"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Install required packages in Colab\n",
    "!pip install git+https://github.com/IBM/terratorch.git\n",
    "!pip install rioxarray matplotlib\n",
    "!wget https://raw.githubusercontent.com/IBM/terramind/refs/heads/main/notebooks/plotting_utils.py"
   ],
   "id": "4e358e3488e2e032"
  },
  {
   "metadata": {
    "collapsed": true
   },
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "import torch\n",
    "import rioxarray as rxr\n",
    "import matplotlib.pyplot as plt\n",
    "from terratorch.registry import FULL_MODEL_REGISTRY\n",
    "from plotting_utils import plot_s2, plot_modality\n",
    "\n",
    "# Select device\n",
    "if torch.cuda.is_available():\n",
    "    device = 'cuda'\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = 'mps'\n",
    "else:\n",
    "    device = 'cpu'"
   ],
   "id": "6ead3d84cdb0cdc4"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Build model that generates S-2 L2A (via tokenizer), S-1 GRD, S-1 RTC, DEM, LULC, and NDVI based on raw S-2 L2A input.\n",
    "# Please note that each output modality requires its own tokenizer which increases the memory requirements\n",
    "model = FULL_MODEL_REGISTRY.build(\n",
    "    'terramind_v1_base_generate',\n",
    "    modalities=['S2L2A'],  # Define the input\n",
    "    output_modalities=['S1GRD', 'DEM', 'LULC'],  # Define the output from S2L2A, S1GRD, S1RTC, DEM, LULC, and NDVI\n",
    "    pretrained=True,\n",
    "    standardize=True,  # If standardize=True, you don't need to do the standardization yourself.\n",
    "    # offset={'S2L2A': 1000}  # Optional offset in your data. The offset is also applied to the generation. \n",
    ")\n",
    "\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "\n",
    "# If you need the standardization values, see\n",
    "from terratorch.models.backbones.terramind.model.terramind_register import v1_pretraining_mean, v1_pretraining_std"
   ],
   "id": "72ebcbdd51a967ff",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "if os.path.exists('../examples'):\n",
    "  # Load local S-2 L2A example paths\n",
    "  examples = [\n",
    "      '../examples/S2L2A/38D_378R_2_3.tif',\n",
    "      '../examples/S2L2A/282D_485L_3_3.tif',\n",
    "      '../examples/S2L2A/433D_629L_3_1.tif',\n",
    "      '../examples/S2L2A/637U_59R_1_3.tif',\n",
    "      '../examples/S2L2A/609U_541L_3_0.tif',\n",
    "  ]\n",
    "else:\n",
    "  # Download examples to Colab\n",
    "  !wget -q https://raw.githubusercontent.com/IBM/terramind/refs/heads/main/examples/S2L2A/38D_378R_2_3.tif\n",
    "  !wget -q https://raw.githubusercontent.com/IBM/terramind/refs/heads/main/examples/S2L2A/282D_485L_3_3.tif\n",
    "  !wget -q https://raw.githubusercontent.com/IBM/terramind/refs/heads/main/examples/S2L2A/433D_629L_3_1.tif\n",
    "  !wget -q https://raw.githubusercontent.com/IBM/terramind/refs/heads/main/examples/S2L2A/637U_59R_1_3.tif\n",
    "  !wget -q https://raw.githubusercontent.com/IBM/terramind/refs/heads/main/examples/S2L2A/541L_3_0.tif\n",
    "  examples = [\n",
    "      '38D_378R_2_3.tif',\n",
    "      '282D_485L_3_3.tif',\n",
    "      '433D_629L_3_1.tif',\n",
    "      '637U_59R_1_3.tif',\n",
    "      '609U_541L_3_0.tif',\n",
    "  ]"
   ],
   "id": "804d6b8fdeebd9f3"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "example_id = 1  # Select id between 0 and 4\n",
    "# Load an S-2 L2A example\n",
    "data = rxr.open_rasterio(examples[example_id])\n",
    "# Convert to shape [B, C, 224, 224]\n",
    "data = torch.Tensor(data.values, device='cpu').unsqueeze(0)"
   ],
   "id": "f399d4fb83a5adfa",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Visualize S-2 L2A input as RGB\n",
    "plot_s2(data)"
   ],
   "id": "93f674b55d92ab9",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Run model with diffusion steps\n",
    "input = data.to(device)\n",
    "with torch.no_grad():\n",
    "  generated = model(input, verbose=True, timesteps=10)"
   ],
   "id": "7d9754c331e226cf",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Plot generations\n",
    "n_plots = len(generated) + 1\n",
    "fig, ax = plt.subplots(1, n_plots, figsize=(5 * n_plots, 5))\n",
    "\n",
    "plot_s2(input, ax=ax[0])\n",
    "ax[0].set_title('Input')\n",
    "\n",
    "for i, (mod, value) in enumerate(generated.items()):\n",
    "    plot_modality(mod, value, ax=ax[i + 1])\n",
    "\n",
    "    ax[i+1].set_title('generated ' + mod)\n",
    "    \n",
    "plt.show()"
   ],
   "id": "baf5b31a642779e0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "f0e4c84c77ca50db",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}