Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "###Import required packages###\n",
    "# Make sure you have all the packages listed below\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"/usr/local/lib/python2.7/site-packages/\")#path to OpenCV\n",
    "import os\n",
    "import argparse\n",
    "import cv2\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "##Code block to generate images corresponding to each character\n",
    "#Character images are sorted in the order in wich they appear in the image.\n",
    "\n",
    "# Load the image\n",
    "img = cv2.imread('temp/1-this_is_a_test-1.png')\n",
    "\n",
    "# convert to grayscale\n",
    "gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "# Apply adaptive threshold ()\n",
    "thresh = cv2.adaptiveThreshold(gray,255,1,1,11,2)\n",
    "\n",
    "# Apply some dilation and erosion to join the gaps\n",
    "thresh = cv2.dilate(thresh,None,iterations = 8)\n",
    "thresh = cv2.erode(thresh,None,iterations = 8)\n",
    "\n",
    "# Find the contours for each WORD\n",
    "contours,hierarchy = cv2.findContours(thresh,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "# Create a folder to store the output files for segmentation\n",
    "output_folder = \"temp/segmentation_pdfimage/\"\n",
    "if(not os.path.exists(output_folder)):\n",
    "    os.mkdir(output_folder)\n",
    "\n",
    "word_count = 0\n",
    "\n",
    "#Sort the contours for the words\n",
    "boundingBoxes = [cv2.boundingRect(c) for c in contours]\n",
    "(contours, boundingBoxes) = zip(*sorted(zip(contours, boundingBoxes),key=lambda b:b[1][0], reverse=False))\n",
    "\n",
    "#Get the characters from inside each word using word's contour\n",
    "for cnt in contours:\n",
    "    #Get the word using the original image\n",
    "    x,y,w,h = cv2.boundingRect(cnt)\n",
    "    word_img = img[y:y+h]\n",
    "    word_img = word_img[:,x:x+w]\n",
    "    \n",
    "    #Convert it to grayscale\n",
    "    word_gray = cv2.cvtColor(word_img,cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    #apply thresholding\n",
    "    word_thresh = cv2.adaptiveThreshold(word_gray,255,1,1,11,2)\n",
    "\n",
    "    #Find contours inside the word, each contour will represent a character\n",
    "    word_contours,word_hierarchy = cv2.findContours(word_thresh,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)\n",
    "    boundingBoxes = [cv2.boundingRect(c) for c in word_contours]\n",
    "    \n",
    "    #Sort the contours for each word\n",
    "    (word_contours, boundingBoxes) = zip(*sorted(zip(word_contours, boundingBoxes),key=lambda b:b[1][0], reverse=False))\n",
    "\n",
    "    char_count = 0\n",
    "    \n",
    "    #Generate image file corresponding to each character in the word\n",
    "    for char_cnt in word_contours:\n",
    "        char_x,char_y, char_w,char_h = cv2.boundingRect(char_cnt)\n",
    "        char = word_img\n",
    "        char = char[:,char_x:char_x+char_w]\n",
    "        height = len(word_img)\n",
    "        if (char_h < height/2) : #remove noise\n",
    "            continue;\n",
    "        #position the image in center\n",
    "        if( height < char_w ):\n",
    "            top_border = (char_w - height)/2;\n",
    "            bottom_border = (char_w - height - top_border)\n",
    "            char= cv2.copyMakeBorder(char,top_border,bottom_border,0,0,cv2.BORDER_CONSTANT,value=[255,255,255])\n",
    "\n",
    "        if( char_w < height ):\n",
    "            left_border = (height - char_w)/2\n",
    "            right_border = (height - char_w - left_border)\n",
    "            char= cv2.copyMakeBorder(char,0,0,left_border,right_border,cv2.BORDER_CONSTANT,value=[255,255,255])\n",
    "\n",
    "        #Our caffe model was trained on 128*128, so we resize the image to the same size    \n",
    "        char = cv2.resize(char, (128, 128))\n",
    "        cv2.imwrite(output_folder+\"word_\"+str(word_count)+\"char_\"+str(char_count)+\".png\",char)\n",
    "        char_count +=1\n",
    "    word_count += 1"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}