theia.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Theia"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gather intelligence from satellite data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check out the [intelligenerator organization](https://github.com/intelligenerator) on Github for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unet_dataset import SatelliteImageDataset\n",
    "from unet import UNet\n",
    "\n",
    "import src.data as data_utils "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 1\n",
    "EPOCHS = 10\n",
    "MODEL_PATH = './model/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = SatelliteImageDataset(images_dir='data/images', targets_dir='data/targets', transform=data_utils.transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, collate_fn=data_utils.collate\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_img(img, transpose=True):\n",
    "    # img = img / 2 + 0.5 # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    if transpose:\n",
    "        npimg = npimg.transpose(1, 2, 0)\n",
    "    plt.imshow(npimg)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataiter = iter(dataloader)\n",
    "imgs, labels = next(dataiter)\n",
    "\n",
    "y = labels[0]\n",
    "print(y.shape)\n",
    "print(y.unique())\n",
    "\n",
    "show_img(imgs[0])\n",
    "show_img(y.squeeze(), transpose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device \"%s\" for training' % dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = UNet(in_channels=3, out_channels=1, padding=True).to(dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(net.parameters(), lr=0.0001)\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "loss_history = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "for epoch in range(EPOCHS):\n",
    "    running_loss = 0.0\n",
    "    \n",
    "    for i, (imgs, labels) in enumerate(dataloader):        \n",
    "        imgs = imgs.to(dev)\n",
    "        labels = labels.to(dev)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        y = net(imgs)\n",
    "        \n",
    "        loss = criterion(y, labels)\n",
    "        loss.backward()\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "        \n",
    "        if i % 50 == 49:\n",
    "            running_loss /= 50\n",
    "            \n",
    "            print('[%.3d / %.3d] Loss: %.9f' % (epoch, i, running_loss))\n",
    "            loss_history.append(running_loss)\n",
    "            \n",
    "            running_loss = 0\n",
    "        \n",
    "    torch.save(net.state_dict(), MODEL_PATH + 'theia.pth')\n",
    "    with open(MODEL_PATH + 'history.json', 'w+') as f:\n",
    "        json.dump(loss_history, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}