Dataset Creator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This Notebook is sequentially processing Images from SuperCaustics to be used in Cleargrasp Training. \n",
    "The Code currently does not check for many things that can go wrong during copy-pasting data, so definitely make backups from your images before running this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import imageio\n",
    "import PIL.Image as Image\n",
    "from PIL import ImageFilter\n",
    "from PIL import ImageChops\n",
    "import io\n",
    "import os\n",
    "from glob import glob\n",
    "import shutil\n",
    "import cv2\n",
    "import random as r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageio.plugins.freeimage.download()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make Dataset Folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#point this to the RAW Screenshots folder. Make Backups before running. You can Comment Out any files that you are not collecting.\n",
    "os.chdir('E:/')\n",
    "dir_src = ('E:/SuperCaustics/Saved/Screenshots/Windows/')\n",
    "\n",
    "Depth =   (dir_src + \"Depth/\")\n",
    "Normal = (dir_src + \"/Normal/\")\n",
    "Pic =     (dir_src + \"/Picture_No_Caustic/\")\n",
    "Caustic = (dir_src + \"/Picture_Caustic/\")\n",
    "Caustic_Seg = (dir_src + \"/Caustic_Seg/\")\n",
    "Outlines = (dir_src + \"/Outlines/\")\n",
    "Outlines_Processed = (dir_src + \"/Outlines_Processed/\")\n",
    "Seg = (dir_src + \"/Seg/\")\n",
    "Seg_Processed = (dir_src + \"/Seg_Processed/\")\n",
    "Depth_exr =   (dir_src + \"/Depth_exr/\")\n",
    "Normal_exr = (dir_src + \"/Normal_exr/\")\n",
    "Seg_exr = (dir_src + \"/Seg_exr/\")\n",
    "\n",
    "try:\n",
    "    os.mkdir(Depth)\n",
    "    os.mkdir(Normal)\n",
    "    os.mkdir(Pic)\n",
    "    os.mkdir(Caustic_Seg)\n",
    "    os.mkdir(Caustic)\n",
    "    os.mkdir(Outlines)\n",
    "    os.mkdir(Seg)\n",
    "    os.mkdir(Outlines_Processed)\n",
    "    os.mkdir(Seg_Processed)\n",
    "    os.mkdir(Depth_exr)\n",
    "    os.mkdir(Normal_exr)\n",
    "    os.mkdir(Seg_exr)\n",
    "\n",
    "except OSError:\n",
    "        print('Directory not created.')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Move images into their respective folders (Order is important)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the Same order as you did when capturing your data from AIP SuperCaustics. In my case, I did this in the following order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "caustic = 1\n",
    "counter = 0\n",
    "for filename in os.listdir(dir_src):\n",
    "\n",
    "    if counter == 0:\n",
    "        if filename.endswith('.png'):\n",
    "            shutil.move(dir_src + filename, Caustic) #Caustic\n",
    "            print(\"moved \" + filename, end='\\r')\n",
    "            counter = counter+1\n",
    "\n",
    "    elif counter == 1:\n",
    "        if filename.endswith('.png'):\n",
    "            shutil.move(dir_src + filename, Pic) #Pic\n",
    "            print(\"moved \" + filename, end='\\r')\n",
    "            counter = counter+1\n",
    "            \n",
    "    elif counter == 2:\n",
    "        if filename.endswith('.png'):\n",
    "            shutil.move(dir_src + filename, Depth) #Depth\n",
    "            print(\"moved \" + filename, end='\\r')\n",
    "            counter = 3\n",
    "            \n",
    "    elif counter == 3:\n",
    "        if filename.endswith('.png'):\n",
    "            shutil.move(dir_src + filename, Seg) #Seg\n",
    "            print(\"moved \" + filename, end='\\r')\n",
    "            counter = 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rename raw images in each folder to have the same name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "i = 0\n",
    "for filename in os.listdir(Pic):\n",
    "    os.rename(Pic + '/' + filename, Pic + '/' + str(i) + '.png')\n",
    "    i = i + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n",
    "    \n",
    "v = 0\n",
    "for filename in os.listdir(Caustic):\n",
    "    os.rename(Caustic + '/' + filename, Caustic + '/' + str(v) + '.png')\n",
    "    v = v + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n",
    "\n",
    "j = 0\n",
    "for filename in os.listdir(Depth):\n",
    "    os.rename(Depth + '/' + filename, Depth + '/' + str(j) + '.png')\n",
    "    j = j + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n",
    "\n",
    "k = 0\n",
    "for filename in os.listdir(Normal):\n",
    "    os.rename(Normal + '/' + filename, Normal + '/' + str(k) + '.png')\n",
    "    k = k + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n",
    "    \n",
    "l = 0\n",
    "for filename in os.listdir(Outlines):\n",
    "    os.rename(Outlines + '/' + filename, Outlines + '/' + str(l) + '.png')\n",
    "    l = l + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n",
    "    \n",
    "m = 0\n",
    "for filename in os.listdir(Seg):\n",
    "    os.rename(Seg + '/' + filename, Seg + '/' + str(m) + '.png')\n",
    "    m = m + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n",
    "    \n",
    "n = 0\n",
    "for filename in os.listdir(Depth):\n",
    "    os.rename(Depth + '/' + filename, Depth + '/' + str(n) + '.png')\n",
    "    n = n + 1\n",
    "    print(\"renamed \" + filename, end='\\r')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seg & Outlines should be L."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for filename in os.listdir(Seg):\n",
    "    segaddress = Seg + filename\n",
    "    newaddress = Seg_Processed + filename\n",
    "    img = []\n",
    "    arr = np.array(Image.open(segaddress).convert('L'))\n",
    "    img = np.where(arr>= 220, 255,-1)\n",
    "    cv2.imwrite(newaddress, img)\n",
    "    print(filename + 'Processed.', end='\\r')\n",
    "    \n",
    "print('Seg Files Processed.')\n",
    "\n",
    "for filename in os.listdir(Outlines):\n",
    "    outaddress = Outlines + filename\n",
    "    newaddress = Outlines_Processed + filename\n",
    "    image = Image.open(outaddress)\n",
    "    image = image.convert('L')\n",
    "    arr = np.asarray(image)\n",
    "    arr = np.floor(arr/126)\n",
    "    cv2.imwrite(newaddress, arr)\n",
    "    print(filename + 'Processed.', end='\\r')\n",
    "    \n",
    "print('Outline Files Processed.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Processing Caustic Segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def CausticGT(image, caustic , seg, thresh = 10):\n",
    "    import PIL.Image as Image\n",
    "    #Binary Function\n",
    "    fn = lambda x : 255 if x > thresh else 0\n",
    "    \n",
    "    img = Image.open(image).convert('RGB')\n",
    "    img_no_caustic = Image.open(caustic).convert('RGB')\n",
    "    segm = Image.open(seg).convert('1')\n",
    "    \n",
    "    #Subtract the Two input Images\n",
    "    Precursor = ImageChops.subtract(img, img_no_caustic)\n",
    "    \n",
    "    #Make Binary Image\n",
    "    r = Precursor.convert('L').point(fn, mode='1')\n",
    "    \n",
    "    #Remove Noise\n",
    "    p = ImageChops.subtract(r, segm)\n",
    "    p = p.filter(ImageFilter.MinFilter(5))\n",
    "    p = p.filter(ImageFilter.MaxFilter(5))\n",
    "    \n",
    "    return p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for filename in os.listdir(Seg_Processed):\n",
    "    Seg = str(Seg_Processed + filename)\n",
    "    Image = str(Caustic + filename)\n",
    "    NoCaustic = str(Pic + filename)\n",
    "    CausticGT(Image, NoCaustic, Seg).save(Caustic_Seg + filename)\n",
    "    print(filename + 'Processed.', end='\\r')"
   ]
  },
 
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Choosing a Random Validation Set: Making the Folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_val = ('E:/Dataset/Dataset/ExampleData_val/')\n",
    "\n",
    "valDepth =   (dir_val + \"/Depth/\")\n",
    "valNormal = (dir_val + \"/Normal/\")\n",
    "valPic =     (dir_val + \"/Picture_No_Caustic/\")\n",
    "valCaustic = (dir_val + \"/Picture_Caustic/\")\n",
    "valOutlines = (dir_val + \"/Outlines/\")\n",
    "valOutlines_Processed = (dir_val + \"/Outlines_Processed/\")\n",
    "valSeg = (dir_val + \"/Seg/\")\n",
    "valSeg_Processed = (dir_val + \"/Seg_Processed/\")\n",
    "valDepth_exr =   (dir_val + \"/Depth_exr/\")\n",
    "valNormal_exr = (dir_val + \"/Normal_exr/\")\n",
    "valSeg_exr = (dir_val + \"/Seg_exr/\")\n",
    "\n",
    "try:\n",
    "    os.mkdir(dir_val)\n",
    "    os.mkdir(valDepth)\n",
    "    os.mkdir(valNormal)\n",
    "    os.mkdir(valPic)\n",
    "    os.mkdir(valCaustic)\n",
    "    os.mkdir(valOutlines)\n",
    "    os.mkdir(valSeg)\n",
    "    os.mkdir(valOutlines_Processed)\n",
    "    os.mkdir(valSeg_Processed)\n",
    "    os.mkdir(valDepth_exr)\n",
    "    os.mkdir(valNormal_exr)\n",
    "    os.mkdir(valSeg_exr)\n",
    "\n",
    "except OSError:\n",
    "        print('Directory not created.')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Moving The Processed Files to Validation Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valcount = 0\n",
    "validationSetSize = 6\n",
    "for filename in os.listdir(Pic):\n",
    "    if valcount < validationSetSize and filename.endswith('.png') and r.randint(0,6) == 6:\n",
    "            \n",
    "            shutil.move(Depth + filename, valDepth)\n",
    "            shutil.move(Normal + filename, valNormal)\n",
    "            shutil.move(Pic + filename, valPic)\n",
    "            shutil.move(Caustic + filename, valCaustic)\n",
    "            shutil.move(Outlines + filename, valOutlines)\n",
    "            shutil.move(Outlines_Processed + filename, valOutlines_Processed)\n",
    "            shutil.move(Seg + filename, valSeg)\n",
    "            shutil.move(Seg_Processed + filename, valSeg_Processed)\n",
    "            shutil.move(Depth_exr + filename.replace(\".png\", \".exr\"), valDepth_exr)\n",
    "            shutil.move(Normal_exr + filename.replace(\".png\", \".exr\"), valNormal_exr)\n",
    "            shutil.move(Seg_exr + filename.replace(\".png\", \".exr\"), valSeg_exr)\n",
    "            \n",
    "            print(\"Chosen \" + filename, end='\\r')\n",
    "            valcount += 1 "
   ]
  }
 ],
 "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
}