-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
e48d1ca
commit 5fee2f5
Showing
12 changed files
with
2,356 additions
and
1,385 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,194 @@ | ||
| { | ||
| "cells": [ | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "import os\n", | ||
| "import tempfile\n", | ||
| "from pathlib import Path\n", | ||
| "\n", | ||
| "from torch.utils.data import DataLoader\n", | ||
| "from torchvision.models.segmentation import fcn_resnet50\n", | ||
| "import torch.nn as nn\n", | ||
| "from torchgeo.datasets import NAIP, ChesapeakeDE, stack_samples\n", | ||
| "from torchgeo.datasets.utils import download_url\n", | ||
| "from torchgeo.samplers import RandomGeoSampler\n", | ||
| "from torch.nn import CrossEntropyLoss\n", | ||
| "from torch.optim import Adam\n", | ||
| "import torch\n", | ||
| "import numpy as np\n", | ||
| "import matplotlib.pyplot as plt" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "from minerva.models import FCN8ResNet18\n", | ||
| "from minerva.utils.utils import get_cuda_device\n", | ||
| "\n", | ||
| "device = get_cuda_device(0)" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "data_root = tempfile.gettempdir()\n", | ||
| "train_root = Path(data_root, \"naip\", \"train\")\n", | ||
| "test_root = Path(data_root, \"naip\", \"test\")\n", | ||
| "naip_url = \"https://naipeuwest.blob.core.windows.net/naip/v002/de/2018/de_060cm_2018/38075/\"\n", | ||
| "tiles = [\n", | ||
| " \"m_3807511_ne_18_060_20181104.tif\",\n", | ||
| " \"m_3807511_se_18_060_20181104.tif\",\n", | ||
| " \"m_3807512_nw_18_060_20180815.tif\",\n", | ||
| "]\n", | ||
| "\n", | ||
| "for tile in tiles:\n", | ||
| " download_url(naip_url + tile, train_root)\n", | ||
| "\n", | ||
| "download_url(naip_url + \"m_3807512_sw_18_060_20180815.tif\", test_root)" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "train_naip = NAIP(train_root)\n", | ||
| "test_naip = NAIP(test_root)\n", | ||
| "\n", | ||
| "chesapeake_root = os.path.join(data_root, \"chesapeake\")\n", | ||
| "\n", | ||
| "chesapeake = ChesapeakeDE(chesapeake_root, crs=train_naip.crs, res=train_naip.res, download=True)\n", | ||
| "\n", | ||
| "train_dataset = train_naip & chesapeake\n", | ||
| "test_dataset = test_naip & chesapeake\n", | ||
| "\n", | ||
| "sampler = RandomGeoSampler(train_naip, size=256, length=200)\n", | ||
| "dataloader = DataLoader(train_dataset, sampler=sampler, collate_fn=stack_samples, batch_size=32)\n", | ||
| "\n", | ||
| "testsampler = RandomGeoSampler(test_naip, size=256, length=8)\n", | ||
| "testdataloader = DataLoader(test_dataset, sampler=testsampler, collate_fn=stack_samples, batch_size=8, num_workers=4)\n", | ||
| "testdata = list(testdataloader)[0]" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "crit = CrossEntropyLoss()\n", | ||
| "\n", | ||
| "# Criterions are normally parsed to models at init in minerva.\n", | ||
| "fcn = FCN8ResNet18(crit, input_size=(4, 256, 256), n_classes=13).to(device)\n", | ||
| "opt = Adam(fcn.parameters(), lr=1e-3)\n", | ||
| "\n", | ||
| "# Optimisers need to be set to a model in minerva before training.\n", | ||
| "fcn.set_optimiser(opt)\n", | ||
| "\n", | ||
| "for epoch in range(101):\n", | ||
| " losses = []\n", | ||
| " for i, sample in enumerate(dataloader):\n", | ||
| " image = sample[\"image\"].to(device).float() / 255.0\n", | ||
| " target = sample[\"mask\"].to(device).long().squeeze(1)\n", | ||
| " \n", | ||
| " # Uses MinervaModel.step.\n", | ||
| " loss, pred = fcn.step(image, target, train=True)\n", | ||
| " losses.append(loss.item())\n", | ||
| "\n", | ||
| " print(epoch, np.mean(losses))\n", | ||
| " if epoch % 10 == 0:\n", | ||
| " with torch.no_grad():\n", | ||
| " image = testdata[\"image\"].to(device).float() / 255.0\n", | ||
| " target = testdata[\"mask\"].to(device).long().squeeze(1)\n", | ||
| " pred = fcn(image)\n", | ||
| "\n", | ||
| " fig, axs = plt.subplots(3, pred.shape[0], figsize=(10,4))\n", | ||
| " for i in range(pred.shape[0]):\n", | ||
| " axs[0,i].imshow(image[i].cpu().numpy()[:3].transpose(1,2,0))\n", | ||
| " axs[1,i].imshow(target[i].cpu().numpy(), cmap=\"Set3\", vmin=0, vmax=12)\n", | ||
| " axs[2,i].imshow(pred[i].detach().argmax(dim=0).cpu().numpy(), cmap=\"Set3\", vmin=0, vmax=12)\n", | ||
| " plt.setp(plt.gcf().get_axes(), xticks=[], yticks=[])\n", | ||
| " plt.show()" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "fcn = fcn_resnet50(num_classes=13).to(device)\n", | ||
| "fcn.backbone.conv1 = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False).to(device)\n", | ||
| "\n", | ||
| "crit = CrossEntropyLoss()\n", | ||
| "opt = Adam(fcn.parameters(), lr=1e-3)\n", | ||
| "\n", | ||
| "for epoch in range(101):\n", | ||
| " losses = []\n", | ||
| " for i, sample in enumerate(dataloader):\n", | ||
| " image = sample[\"image\"].to(device).float() / 255.0\n", | ||
| " target = sample[\"mask\"].to(device).long().squeeze(1)\n", | ||
| "\n", | ||
| " opt.zero_grad()\n", | ||
| " pred = fcn(image)[\"out\"]\n", | ||
| " loss = crit(pred, target)\n", | ||
| " loss.backward()\n", | ||
| " opt.step()\n", | ||
| " losses.append(loss.item())\n", | ||
| "\n", | ||
| " print(epoch, np.mean(losses))\n", | ||
| " if epoch % 10 == 0:\n", | ||
| " with torch.no_grad():\n", | ||
| " image = testdata[\"image\"].to(device).float() / 255.0\n", | ||
| " target = testdata[\"mask\"].to(device).long().squeeze(1)\n", | ||
| " pred = fcn(image)[\"out\"]\n", | ||
| "\n", | ||
| " fig, axs = plt.subplots(3, pred.shape[0], figsize=(10,4))\n", | ||
| " for i in range(pred.shape[0]):\n", | ||
| " axs[0,i].imshow(image[i].cpu().numpy()[:3].transpose(1,2,0))\n", | ||
| " axs[1,i].imshow(target[i].cpu().numpy(), cmap=\"Set3\", vmin=0, vmax=12)\n", | ||
| " axs[2,i].imshow(pred[i].detach().argmax(dim=0).cpu().numpy(), cmap=\"Set3\", vmin=0, vmax=12)\n", | ||
| " plt.setp(plt.gcf().get_axes(), xticks=[], yticks=[])\n", | ||
| " plt.show()" | ||
| ] | ||
| } | ||
| ], | ||
| "metadata": { | ||
| "kernelspec": { | ||
| "display_name": "minerva-310", | ||
| "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.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0]" | ||
| }, | ||
| "orig_nbformat": 4, | ||
| "vscode": { | ||
| "interpreter": { | ||
| "hash": "3564bae54b830248e5fcf548a4e349b732e585ece6f047dc1ae97c29756580ff" | ||
| } | ||
| } | ||
| }, | ||
| "nbformat": 4, | ||
| "nbformat_minor": 2 | ||
| } |
Oops, something went wrong.