create_traintestsplit.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "import os\n",
    "import cv2\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "directory = 'aug_data/train'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#edge images\n",
    "edge_X = []\n",
    "edge_Y = []\n",
    "for name in os.listdir(directory+'/edge'):\n",
    "    if name != '_DS_Store':\n",
    "        img = cv2.imread(directory+'/edge/'+name,0)\n",
    "        edge_X.append(img)\n",
    "        edge_Y.append([0,1,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(250, 598, 434)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_X = np.asarray(edge_X)\n",
    "edge_X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(250, 3)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_Y = np.asarray(edge_Y)\n",
    "edge_Y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_X_train, edge_X_test, edge_y_train, edge_y_test = train_test_split(edge_X, edge_Y, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#edge images\n",
    "good_X = []\n",
    "good_Y = []\n",
    "for name in os.listdir(directory+'/good'):\n",
    "    if name != '_DS_Store':\n",
    "        img = cv2.imread(directory+'/good/'+name,0)\n",
    "        good_X.append(img)\n",
    "        good_Y.append([1,0,0])\n",
    "        \n",
    "good_X = np.asarray(good_X)\n",
    "good_Y = np.asarray(good_Y)\n",
    "good_X_train, good_X_test, good_y_train, good_y_test = train_test_split(good_X, good_Y, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#edge images\n",
    "water_X = []\n",
    "water_Y = []\n",
    "for name in os.listdir(directory+'/water'):\n",
    "    if name != '_DS_Store':\n",
    "        img = cv2.imread(directory+'/water/'+name,0)\n",
    "        water_X.append(img)\n",
    "        water_Y.append([0,0,1])\n",
    "        \n",
    "water_X = np.asarray(water_X)\n",
    "water_Y = np.asarray(water_Y)\n",
    "water_X_train, water_X_test, water_y_train, water_y_test = train_test_split(water_X, water_Y, test_size=0.05, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = []\n",
    "test_X = []\n",
    "train_y = []\n",
    "test_y = []\n",
    "for i in good_X_train:\n",
    "    train_X.append(i)\n",
    "for i in good_y_train:\n",
    "    train_y.append(i)\n",
    "for i in edge_X_train:\n",
    "    train_X.append(i)\n",
    "for i in edge_y_train:\n",
    "    train_y.append(i)\n",
    "for i in water_X_train:\n",
    "    train_X.append(i)\n",
    "for i in water_y_train:\n",
    "    train_y.append(i)\n",
    "    \n",
    "for i in good_X_test:\n",
    "    test_X.append(i)\n",
    "for i in good_y_test:\n",
    "    test_y.append(i)\n",
    "for i in edge_X_test:\n",
    "    test_X.append(i)\n",
    "for i in edge_y_test:\n",
    "    test_y.append(i)\n",
    "for i in water_X_test:\n",
    "    test_X.append(i)\n",
    "for i in water_y_test:\n",
    "    test_y.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = np.asarray(train_X)\n",
    "test_X = np.asarray(test_X)\n",
    "train_y = np.asarray(train_y)\n",
    "test_y = np.asarray(test_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(598, 598, 434)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "251\n",
      "200\n",
      "250\n"
     ]
    }
   ],
   "source": [
    "print(len(water_X))\n",
    "print(len(good_X))\n",
    "print(len(edge_X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(103, 598, 434)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = {\n",
    "    'train_X':train_X,\n",
    "    'train_y':train_y,\n",
    "    'test_X':test_X,\n",
    "    'test_y':test_y\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"msft.pickle\", \"wb\") as output_file:\n",
    "    pickle.dump(d, output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}