MiniBatch_generator_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import csv\n",
    "import os\n",
    "import torch\n",
    "\n",
    "import random\n",
    "\n",
    "from PIL import Image\n",
    "from skimage import io\n",
    "\n",
    "verbose = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class mini_batch(torch.nn.Module):\n",
    "    def __init__(self, dat_folder,transform,percent):\n",
    "        super(mini_batch, self).__init__()\n",
    "        self.root_dir = os.path.dirname(dat_folder) \n",
    "        self.transform = transform\n",
    "        self.percent = percent\n",
    "\n",
    "        #list of folders \n",
    "        ind = [individual for individual in os.listdir(self.root_dir) if os.path.isdir(os.path.join(self.root_dir,individual))]\n",
    "            \n",
    "        self.ind = ind        \n",
    "        \n",
    "    def loadfile(self, data_file):\n",
    "        image = io.imread(data_file)\n",
    "        if len(image.shape)<3:\n",
    "            image = np.stack((image,)*3, axis=-1)\n",
    "        return image\n",
    "    \n",
    "    \n",
    "    def prepare_batch(self):\n",
    "\n",
    "        train_line_list = []\n",
    "        test_line_list = []\n",
    "        \n",
    "        labels_train = []\n",
    "        labels_test = []\n",
    "\n",
    "        for i,individual in enumerate(self.ind):\n",
    "            folder_path = os.path.join(self.root_dir,individual)\n",
    "            name_images = [name for name in os.listdir(folder_path)]  \n",
    "\n",
    "            random.shuffle(name_images)\n",
    "            N = len(name_images)\n",
    "\n",
    "            for j in range(int(N*self.percent)):\n",
    "                if os.path.isfile(os.path.join(folder_path,name_images[j])):\n",
    "                    image_filename = os.path.join(individual,name_images[j])\n",
    "                    if verbose: print(\"image {}{}\".format(j,image_filename))\n",
    "                    train_line_list.append(image_filename)\n",
    "                    labels_train.append(i)\n",
    " \n",
    "\n",
    "            for j in range(int(N*self.percent),N):\n",
    "                if os.path.isfile(os.path.join(folder_path,name_images[j])):\n",
    "                    image_filename = os.path.join(individual,name_images[j])\n",
    "                    if verbose: print(\"image {}{}\".format(j,image_filename))\n",
    "                    test_line_list.append(image_filename)\n",
    "                    labels_test.append(i)\n",
    "\n",
    "        image_train = [self.transform(Image.fromarray(self.loadfile(os.path.join(self.root_dir,train_line_list[i])))) for i in range(len(train_line_list))]                    \n",
    "        image_test = [self.transform(Image.fromarray(self.loadfile(os.path.join(self.root_dir,test_line_list[i])))) for i in range(len(test_line_list))]                    \n",
    "        return image_train, labels_train, image_test, labels_test\n",
    "    \n",
    "    "
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}