Data_Retriever.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "latest-crawford",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "import pandas as pd\n",
    "import os\n",
    "import cv2\n",
    "from PIL import Image\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import transforms\n",
    "from albumentations import (HorizontalFlip, Normalize, Compose, Resize)\n",
    "from albumentations.pytorch import ToTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "funky-aurora",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageData(Dataset):\n",
    "    def __init__(self, df, data_folder, mean, std, phase):\n",
    "        self.df = df\n",
    "        self.root = data_folder\n",
    "        self.mean = mean\n",
    "        self.std = std\n",
    "        self.phase = phase\n",
    "        self.transforms = self.get_transforms()\n",
    "        self.fnames = self.df.index.tolist()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.fnames)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        image_id, mask = self.make_mask(idx)\n",
    "        image_path = os.path.join(self.root, \"train_images\",  image_id)\n",
    "        img = cv2.imread(image_path)\n",
    "        augmented = self.transforms(image=img, mask=mask)\n",
    "        img = augmented['image']\n",
    "        mask = augmented['mask'] # 1x256x1600x4\n",
    "        mask = mask[0].permute(2, 0, 1) # 4x256x1600\n",
    "        return img, mask\n",
    "\n",
    "        \n",
    "    def make_mask(self, row_id):\n",
    "        '''Given a row index, return image_id and mask (256, 1600, 4) from the dataframe `df`'''\n",
    "        fname = self.df.iloc[row_id].name\n",
    "        labels = self.df.iloc[row_id][:4]\n",
    "        masks = np.zeros((256, 1600, 4), dtype=np.float32) # float32 is V.Imp\n",
    "        # 4:class 1～4 (ch:0～3)\n",
    "        \n",
    "        for idx, label in enumerate(labels.values):\n",
    "            if type(label) is str:\n",
    "                label = label.split(\" \")\n",
    "                positions = map(int, label[0::2])\n",
    "                length = map(int, label[1::2])\n",
    "                mask = np.zeros(256 * 1600, dtype=np.uint8)\n",
    "                for pos, le in zip(positions, length):\n",
    "                    mask[pos:(pos + le)] = 1\n",
    "                masks[:, :, idx] = mask.reshape(256, 1600, order='F')\n",
    "        return fname, masks\n",
    "    \n",
    "    def get_transforms(self):\n",
    "        list_transforms = []\n",
    "        if self.phase == \"train\":\n",
    "            list_transforms.extend(\n",
    "                [\n",
    "                    HorizontalFlip(p=0.5), # only horizontal flip as of now\n",
    "                ]\n",
    "            )\n",
    "        list_transforms.extend(\n",
    "            [\n",
    "                Normalize(mean=self.mean, std=self.std, p=1),\n",
    "                Resize(224,224),\n",
    "                ToTensor(),\n",
    "            ]\n",
    "        )\n",
    "        list_trfms = Compose(list_transforms)\n",
    "        return list_trfms\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eligible-password",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}