---
# System prepended metadata

title: Untitled

---

import cv2
import pytesseract
from matplotlib import pyplot as plt
import numpy as np
import re
import argparse


parser = argparse.ArgumentParser(description='Take skew angle and image as input and return the rotation angle.')

parser.add_argument('-p', '--path', required = True, help='Path to the input image')
parser.add_argument('-sa', '--skew-angle', help='Skew angle of the image')
parser.add_argument('-t', '--text', action = 'store_true', help='If you want to extract text from the image,\
 you can add -t to the program call')

args = parser.parse_args()

def words():
    words_ = ''
    with open('./words/en') as reader:
        words_ += reader.read()
    return words_.split("\n")

def tokenize(text):
    pat = r'[ /\',\".?!&:;|\(\)\[\]\{\}]+'
    return re.split(pat, text)

def rotate_img(img, angle):
    (h, w) = img.shape[:2]
    (cX, cY) = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    return cv2.warpAffine(img, M, (nW, nH))


def get_rotation_angle(img, possible_rotation_angles):
    rotation_angle = 0
    most_no_of_words = 0
    text = ""
    for angle in possible_rotation_angles:
        rotated_img = rotate_img(img, angle)
        current_text = pytesseract.image_to_string(rotated_img).lower()
        tokens = tokenize(current_text)
        no_of_words = len(set(words()) & set(tokens))
        if no_of_words > most_no_of_words:
            rotation_angle = angle
            most_no_of_words = no_of_words
    return rotation_angle


img = cv2.imread(args.path)
gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
T, binarized = cv2.threshold(gray_img, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY)


if args.skew_angle:
    skew_angle = int(args.skew_angle)
    if skew_angle + 180 > 359:
        possible_rotation_angles = [skew_angle, skew_angle - 180]
    else:
        possible_rotation_angles = [skew_angle, skew_angle + 180]
else:
    possible_rotation_angles = range(0, 359, 10)


rotation_angle = get_rotation_angle(binarized, possible_rotation_angles)
print("This image should be rotated by {} degree to perform OCR".format(rotation_angle))

if args.text:
    rotated = rotate_img(binarized, rotation_angle)
    eroded = cv2.erode(rotated, np.ones((2,2)))
    dilated = cv2.dilate(rotated, np.ones((2,2)))
    er_text = pytesseract.image_to_string(eroded)
    dl_text = pytesseract.image_to_string(dilated)
    text_ = er_text if len(er_text) > len(dl_text) else dl_text
    print(text_)