SR_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import listdir\n",
    "from os.path import join\n",
    "from PIL import Image\n",
    "from torch.utils.data.dataset import Dataset\n",
    "import torchvision.transforms as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_image_file(filename):\n",
    "    return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_valid_crop_size(crop_size, upscale_factor):\n",
    "    return crop_size - (crop_size % upscale_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_HR_transform(crop_size):\n",
    "    return T.Compose([ T.RandomCrop(crop_size, pad_if_needed=True), T.ToTensor() ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_LR_transform(crop_size, upscale_factor):\n",
    "    transform = T.Compose([\n",
    "        T.ToPILImage(),\n",
    "        T.Resize(crop_size // upscale_factor, interpolation = Image.BICUBIC),\n",
    "        T.ToTensor(),\n",
    "    ])\n",
    "    \n",
    "    return transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class dataset_train_from_folder(Dataset):\n",
    "    \n",
    "    \"\"\"\n",
    "    args: \n",
    "    dataset_dir: (str) directory of datset\n",
    "    crop_size: (int) image size of HR\n",
    "    upscale_factor: (int) upscaling factor for processing LR from HR\n",
    "    \n",
    "    returns: \n",
    "    LR_image: tensor of size (crop_size // upscale_factor, crop_size // upscale_factor)\n",
    "    HR_restore_image:  tensor of size (crop_size, crop_size)\n",
    "    HR_image: tensor of size (crop_size, crop_size)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dataset_dir, crop_size, upscale_factor):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.image_filename = [join(dataset_dir, x) for x in listdir(dataset_dir) if is_image_file(x)]\n",
    "        self.crop_size = calculate_valid_crop_size(crop_size, upscale_factor)\n",
    "        self.HR_transform = train_HR_transform(crop_size) \n",
    "        self.LR_transform = train_LR_transform(crop_size, upscale_factor)\n",
    "        \n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        HR_image = self.HR_transform(Image.open(self.image_filename[idx]))\n",
    "        LR_image = self.LR_transform(HR_image)\n",
    "        \n",
    "        return LR_image, HR_image\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.image_filename)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class dataset_val_from_folder(Dataset):\n",
    "    \n",
    "    \"\"\"\n",
    "    args: \n",
    "    dataset_dir: (str) directory of datset\n",
    "    crop_size: (int) image size of HR\n",
    "    upscale_factor: (int) upscaling factor for processing LR from HR\n",
    "    \n",
    "    returns: \n",
    "    LR_image: tensor of size (crop_size // upscale_factor, crop_size // upscale_factor)\n",
    "    HR_restore_image:  tensor of size (crop_size, crop_size)\n",
    "    HR_image: tensor of size (crop_size, crop_size)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dataset_dir, crop_size, upscale_factor):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.image_filenames = [join(dataset_dir, x) for x in listdir(dataset_dir) if is_image_file(x)]\n",
    "        self.upscale_factor = upscale_factor\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        image = Image.open(self.image_filenames[idx])\n",
    "        w, h = image.size\n",
    "        crop_size = calculate_valid_crop_size(min(w,h) , self.upscale_factor)\n",
    "        LR_scale = T.Resize(crop_size // self.upscale_factor, interpolation = Image.BICUBIC)\n",
    "        HR_scale = T.Resize(crop_size, interpolation = Image.BICUBIC)\n",
    "        HR_image = T.CenterCrop(crop_size)(image) \n",
    "        LR_image = LR_scale(HR_image)\n",
    "        HR_restore_image = HR_scale(LR_image)\n",
    "        \n",
    "        return T.ToTensor()(LR_image), T.ToTensor()(HR_restore_image), T.ToTensor()(HR_image)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.image_filenames)\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class dataset_test_from_folder(Dataset):\n",
    "    \n",
    "    \"\"\"\n",
    "    args: \n",
    "    dataset_dir: (str) directory of datset\n",
    "    upscale_factor: (int) upscaling factor for processing LR from HR\n",
    "    \n",
    "    returns: \n",
    "    LR_image: tensor of size (h, w)\n",
    "    HR_restore_image:  tensor of size (h * upscale_factor, w * upscale_factor)\n",
    "    HR_image: tensor of size (h * upscale_factor, w * upscale_factor)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dataset_dir, upscale_factor):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.LR_path = dataset_dir + '/data/'\n",
    "        self.HR_path = dataset_dir + '/target/'\n",
    "        self.upscale_factor = upscale_factor\n",
    "        self.LR_filenames = [join(self.LR_path, x) for x in listdir(self.LR_path) if is_image_file(x)]\n",
    "        self.HR_filenames = [join(self.HR_path, x) for x in listdir(self.HR_path) if is_image_file(x)]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        image_name = self.lr_filenames[index].split('/')[-1]\n",
    "        LR_image = Image.open(self.lr_filenames[index])\n",
    "        w, h = LR_image.size\n",
    "        HR_image = Image.open(self.HR_filenames[index])\n",
    "        HR_scale = T.Resize((self.upscale_factor * h, self.upscale_factor * w), interpolation=Image.BICUBIC)\n",
    "        HR_restore_img = hr_scale(LR_image)\n",
    "        return image_name, ToTensor()(LR_image), ToTensor()(HR_restore_img), ToTensor()(HR_image)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.LR_filenames)"
   ]
  }
 ],
 "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": 2
}