split-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "accompanied-secondary",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import math\n",
    "import random\n",
    "from PIL import Image\n",
    "\n",
    "path = os.getcwd() + \"/frames/img/\"\n",
    "numSamples = 1043"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "blessed-gregory",
   "metadata": {},
   "source": [
    "## Get labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "flying-primary",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = np.empty((numSamples))\n",
    "\n",
    "for s in range(numSamples):\n",
    "    txtFile = open(path+str(s)+\".txt\",\"r\").read()\n",
    "    \n",
    "    if txtFile == \"\":\n",
    "        labels[s] = 0\n",
    "    else:\n",
    "        labels[s] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "double-honey",
   "metadata": {},
   "source": [
    "## Calculate split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "modern-ribbon",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0.0: 503, 1.0: 540}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique, counts = np.unique(labels, return_counts=True)\n",
    "numClassSamples = min(counts) - (min(counts)%4)\n",
    "trainSplit = math.floor(0.7*numClassSamples)\n",
    "testSplit = numClassSamples - trainSplit\n",
    "\n",
    "dict(zip(unique, counts))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "official-baltimore",
   "metadata": {},
   "source": [
    "## Split data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "electoral-waterproof",
   "metadata": {},
   "outputs": [],
   "source": [
    "order = np.arange(numSamples).astype(\"int\")\n",
    "random.shuffle(order)\n",
    "\n",
    "train = np.empty((trainSplit*2))\n",
    "trainBalance = [0,0]\n",
    "test = np.empty((testSplit*2))\n",
    "testBalance = [0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "hydraulic-there",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "500\n",
      "350\n",
      "150\n"
     ]
    }
   ],
   "source": [
    "print(numClassSamples)\n",
    "print(trainSplit)\n",
    "print(testSplit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "noble-palmer",
   "metadata": {},
   "outputs": [],
   "source": [
    "for index in order:\n",
    "    classL = labels[index].astype(\"int\")\n",
    "        \n",
    "    if trainBalance[classL] < trainSplit:\n",
    "        train.put(trainBalance[0] + trainBalance[1], index)\n",
    "        trainBalance[classL] += 1\n",
    "        \n",
    "    elif testBalance[classL] < testSplit:\n",
    "        test.put(testBalance[0] + testBalance[1], index)\n",
    "        testBalance[classL] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "north-carter",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 528.,  635.,  522.,  564.,  356.,  649.,  646.,  911.,  451.,\n",
       "        170.,  240.,  149.,  473.,  888.,  174.,  837.,  359., 1030.,\n",
       "        690.,  172.,  340.,  552.,  924.,  909.,  148.,  368.,  532.,\n",
       "        866.,  420.,  669.,  235.,  711.,  153.,  846.,  233., 1006.,\n",
       "        999.,  710., 1019.,  642.,  229.,  514.,  169.,  833.,  249.,\n",
       "        304.,  739.,  874.,  441.,  445.,  568.,   84.,  210.,  259.,\n",
       "        686.,  727.,  705.,  470.,  483.,  293.,  221.,  720.,   82.,\n",
       "        242.,   32.,  124.,  912.,  219.,  411.,   29.,  628.,  973.,\n",
       "        685., 1026.,  841.,  724.,  853.,  404.,  501.,  196.,   31.,\n",
       "        894.,  486.,   97.,  637.,  971.,  766.,  695.,  314.,  780.,\n",
       "       1041.,  689., 1021.,  180.,  744.,  898.,  777.,  990.,  623.,\n",
       "        191.,   81.,  835.,   93.,  899.,  467.,  165.,  185.,  783.,\n",
       "         14.,  741., 1031.,  722.,  267.,  394.,  832.,  246.,  410.,\n",
       "        491.,  362.,  902.,  155.,  699.,  607.,  464.,  942.,  330.,\n",
       "        481.,  344.,  594.,  488.,  760.,  447.,  923.,  400.,  503.,\n",
       "        523.,  433.,   73.,  838.,  247.,  557.,  756.,  805.,   28.,\n",
       "        900.,   58.,  283.,  784.,  232.,  302.,  487.,  679.,  858.,\n",
       "        319.,  204.,  317.,  622.,  913.,   33.,  228.,  761.,  579.,\n",
       "        586.,  570.,   43.,  408.,  848.,  162.,  428.,   89.,  430.,\n",
       "        826.,  749.,  921.,   54., 1007., 1014.,  129.,  811.,   68.,\n",
       "        755.,  926.,  115.,  746.,  512.,  979.,  790.,  718.,  280.,\n",
       "        895.,  965.,  730.,  694.,  137.,    8.,  791.,  869., 1033.,\n",
       "        496.,  538.,  672.,  333.,   80.,  582.,  484.,   74.,  518.,\n",
       "         39.,  782.,  958.,   66.,  468.,  949., 1002.,  448.,  915.,\n",
       "        353.,  809.,  397.,  536.,  485.,  561.,  950.,  613.,  141.,\n",
       "       1009.,  750.,  286.,   19.,   60.,  403.,  738., 1025.,  995.,\n",
       "        804.,  610.,  732.,  309.,  429.,  814.,  168.,  664.,   48.,\n",
       "         12.,  351.,  880.,  903.,  786.,  700.,  167.,  152.,  972.,\n",
       "        186.,   17.,  719.,  650.,  767.,  845.,  255.,  591.,  994.,\n",
       "        615.,  301.,  215.,  253.,  270.,  959.,  493.,  118.,  960.,\n",
       "        573.,   49.,  944.,  106.,   41.,  287.,  389.,  575.,  792.,\n",
       "          4.,  698.,  206.,  863.,   16.,   21.,  216.,  393.,  214.,\n",
       "        592.,  264.,  776.,  203.,  385.,  772.,   26.,   13.,  946.,\n",
       "        608.,  308.,  943.])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "traditional-thought",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[350, 350]\n",
      "[150, 150]\n"
     ]
    }
   ],
   "source": [
    "print(trainBalance)\n",
    "print(testBalance)\n",
    "train = train.astype(\"int\")\n",
    "test = test.astype(\"int\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "polished-white",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "unique, counts = np.unique(np.concatenate((test,train)), return_counts=True)\n",
    "print(len(counts) == 2*(testSplit+trainSplit))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "industrial-cutting",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "350\n",
      "150\n"
     ]
    }
   ],
   "source": [
    "numTestDoor = 0\n",
    "numTrainDoor = 0\n",
    "\n",
    "for s in range(len(train)):\n",
    "    sample = train[s]\n",
    "    \n",
    "    newPath = os.getcwd() + \"/balancedDataset/train/\" + str(s)\n",
    "    \n",
    "    txtFile = open(path + str(sample) + \".txt\",\"r\").read()\n",
    "    newFile = open(newPath + \".txt\",\"w\")\n",
    "    newFile.write(txtFile)\n",
    "        \n",
    "    image = Image.open(path + str(sample) + \".jpg\")\n",
    "    image = image.save(newPath + \".jpg\")\n",
    "    \n",
    "    if txtFile != \"\":\n",
    "        numTrainDoor += 1\n",
    "\n",
    "for t in range(len(test)):\n",
    "    sample = test[t]\n",
    "    \n",
    "    newPath = os.getcwd() + \"/balancedDataset/test/\" + str(t)\n",
    "    \n",
    "    txtFile = open(path + str(sample) + \".txt\",\"r\").read()\n",
    "    newFile = open(newPath + \".txt\",\"w\")\n",
    "    newFile.write(txtFile)\n",
    "        \n",
    "    image = Image.open(path + str(sample) + \".jpg\")\n",
    "    image = image.save(newPath + \".jpg\")\n",
    "    \n",
    "    if txtFile != \"\":\n",
    "        numTestDoor += 1\n",
    "        \n",
    "print(numTrainDoor)\n",
    "print(numTestDoor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "corporate-attribute",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}