[Inference] Generate pedictions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ad4b8e7a-86e9-45b5-a90f-70292325fdee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import pathlib\n",
    "\n",
    "import tqdm\n",
    "import torch\n",
    "import skimage\n",
    "import rasterio\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from hashtagdeep.models import FCDenseNet103\n",
    "from hashtagdeep.dataset import MiniFranceDFC22\n",
    "from hashtagdeep.utils.colormaps import dfc22_labels_palette"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5398e1f5-8a95-4968-a6e5-f3112c0443d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = MiniFranceDFC22(\n",
    "    base_dir='/home/dubrovin/Projects/Data/DFC2022/',\n",
    "    labeled=False,\n",
    "    val=True,\n",
    "    augmentation=None,\n",
    "    transform=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8afd2121-346e-4932-9013-7ca21b73def1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FCDenseNet103(in_channels=4, n_classes=14)\n",
    "state_dict = torch.load('checkpoints/state_dicts_for_submission_models/experiment00_v1.pt')\n",
    "model.load_state_dict(state_dict)\n",
    "model.eval()\n",
    "model.to('cuda:1');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "aaefecd3-fd7a-4b22-8261-1536c6298184",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████| 2066/2066 [1:05:03<00:00,  1.89s/it]\n"
     ]
    }
   ],
   "source": [
    "base_predictions_dir = pathlib.Path('predictions0')\n",
    "base_predictions_dir.mkdir(exist_ok=True)\n",
    "\n",
    "for i, item in tqdm.tqdm(enumerate(dataset), total=len(dataset)):\n",
    "    x = item['image']\n",
    "    height, width = x.shape[1:]\n",
    "    out = torch.empty(height, width, dtype=torch.int8)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        out[:1000, :1000] = model(x[:, :1000, :1000].unsqueeze(0).to('cuda:1')).cpu().squeeze().argmax(0)\n",
    "        out[1000:, :1000] = model(x[:, 1000:, :1000].unsqueeze(0).to('cuda:1')).cpu().squeeze().argmax(0)\n",
    "        out[:1000, 1000:] = model(x[:, :1000, 1000:].unsqueeze(0).to('cuda:1')).cpu().squeeze().argmax(0)\n",
    "        out[1000:, 1000:] = model(x[:, 1000:, 1000:].unsqueeze(0).to('cuda:1')).cpu().squeeze().argmax(0)\n",
    "    \n",
    "    # transform class indices into class labels\n",
    "    out += 1\n",
    "    \n",
    "    path = dataset.true_color_paths[i]\n",
    "    region, _, file = str(path).split('/')[-3:]\n",
    "    \n",
    "    region_dir = base_predictions_dir / region\n",
    "    region_dir.mkdir(exist_ok=True)\n",
    "    \n",
    "    file_basename = file.split('.')[0]\n",
    "    prediction_file = region_dir / f'{file_basename}_prediction.tif'\n",
    "    \n",
    "    with rasterio.open(path) as src:\n",
    "        transform = src.transform\n",
    "        crs = src.crs\n",
    "    \n",
    "    with rasterio.open(\n",
    "        prediction_file, 'w',\n",
    "        height=height, width=width, count=1, dtype='uint8',\n",
    "        transform=transform, crs=crs,\n",
    "    ) as dst:\n",
    "        dst.write(out[None, :, :])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9 (torch-latest)",
   "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}