notebooks/visualize_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook loads and displays training images with bounding box labels using PyTorch utility functions.\n",
    "\n",
    "Later on, you can enable some albumentations transformations and display the images with transformations applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision.utils import draw_bounding_boxes\n",
    "\n",
    "from wheat.config import load_config\n",
    "from wheat.data_module import WheatDataModule\n",
    "from wheat import visualization as vis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change to the top level directory for this repository\n",
    "os.chdir('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_image_grid(dataset, images_per_row=3, num_rows=3, first_image_index=0):\n",
    "    for irow in range(num_rows):\n",
    "        image_list = []\n",
    "        for icol in range(images_per_row):\n",
    "            image, labels = train_dataset[irow * images_per_row + icol + first_image_index]\n",
    "            result = draw_bounding_boxes(\n",
    "                vis.image_float_to_int_transform(image), labels['boxes'], colors='blue', width=5)\n",
    "            image_list.append(result)\n",
    "        vis.show(image_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the default configuration\n",
    "config = load_config('wheat/config/config.ini')\n",
    "pp = pprint.PrettyPrinter(indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make sure data augmentation is off for first look at training images\n",
    "for key in config['train']['transforms']:\n",
    "    config['train']['transforms'][key] = 0  # set transform probability to 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp.pprint(config['train']['transforms'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the dataset\n",
    "wheat_data_module = WheatDataModule(config)\n",
    "wheat_data_module.setup(stage='fit')\n",
    "train_dataset = wheat_data_module.train_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display a few images\n",
    "plt.rcParams['figure.figsize'] = [15, 10]\n",
    "display_image_grid(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn on one or more transforms\n",
    "config['train']['transforms']['color_jitter_prob'] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# update the transforms associated with the dataset\n",
    "wheat_data_module.config = config\n",
    "train_dataset.transform = wheat_data_module.get_transforms()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display images with the transforms applied\n",
    "display_image_grid(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}