#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# =======================================================
#  tesseractとpytesseractを使って画像ファイルをOCRするスクリプト
#
#  test_pytesseract.py
#  coded by Noboru Harada (noboru@ieee.org)
#
#  Changes:
#  2021/07/06: First version
# =======================================================

import re
import cv2
import numpy as np
import pytesseract
from pytesseract import Output
from matplotlib import pyplot as plt
from operator import itemgetter, attrgetter

# input file
IMG_DIR = './'
filename_img = 'infile.png' 

# text cleansing mode (0: None, 1: remove white space, 2: full cleansing)
mode = 1

###################################
# define preprocessing functions
###################################

# get grayscale image
def get_grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# noise removal
def remove_noise(image):
    return cv2.medianBlur(image,5)
 
#thresholding
def thresholding(image):
    return cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]

#dilation
def dilate(image):
    kernel = np.ones((5,5),np.uint8)
    return cv2.dilate(image, kernel, iterations = 1)
    
#erosion
def erode(image):
    kernel = np.ones((5,5),np.uint8)
    return cv2.erode(image, kernel, iterations = 1)

#opening - erosion followed by dilation
def opening(image):
    kernel = np.ones((5,5),np.uint8)
    return cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)

#canny edge detection
def canny(image):
    return cv2.Canny(image, 100, 200)

#skew correction
def deskew(image):
    coords = np.column_stack(np.where(image > 0))
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
    return rotated

#template matching
def match_template(image, template):
    return cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED) 

#zoom
def resize(image, gain):
    h, w = image.shape[:2]
    inter = cv2.INTER_LINEAR
    #inter = cv2.INTER_NEAREST
    return cv2.resize(image, (w * gain, h * gain), interpolation=inter)

#find squares
def find_squares(image):
    return image

#clean up text
def clean_text(mode, intext):
    text = intext
    if mode == 1: # reduce multiple white spaces into one
        text = re.sub(r'[　 \t]+', ' ', intext)
    elif mode == 2: # cleansing middle spaces in multi-byte characters and remove leading and trailing spaces
        # format reguralization for Japanese char-set
        text = re.sub('([あ-んア-ン一-龥ー、。]) +((?=[あ-んア-ン一-龥ー、。]))', r'\1\2', intext)
        # remove multiple spaces
        text = re.sub(r'[　 \t]+', ' ', text)
        # remove leading and trailing space, 2-byte space and tab space
        text = re.sub(r'^[ 　\t]+', '', text)
        text = re.sub(r'(\| )|([\r\n|\n|\r]) +', r'\1\2', text)
    return text

###################################

# open input image
image = cv2.imread(IMG_DIR + filename_img)

# plot original image
'''
b,g,r = cv2.split(image)
rgb_img = cv2.merge([r,g,b])
plt.figure(figsize=(16,12))
plt.imshow(rgb_img)
plt.title(filename_img)
plt.show()
'''

# Preprocess image
image2 = resize(image, 2)
gray = get_grayscale(image)
thresh = thresholding(gray)
opening = opening(gray)
canny = canny(gray)
images = {'gray': gray, 
          'thresh': thresh, 
          'opening': opening, 
          'canny': canny}

# Plot images after preprocessing
fig = plt.figure(figsize=(20,12))
ax = []
rows = 2
columns = 2
keys = list(images.keys())
for i in range(rows*columns):
    ax.append( fig.add_subplot(rows, columns, i+1) )
    ax[-1].set_title(filename_img + ' - ' + keys[i]) 
    plt.imshow(images[keys[i]], cmap='gray')   

# Get OCR output using Pytesseract
custom_config = r'-l jpn --oem 3 --psm 6'
print('-----------------------------------------')
print('TESSERACT OUTPUT --> ORIGINAL IMAGE')
print('-----------------------------------------')
print(clean_text(mode, pytesseract.image_to_string(image, config=custom_config)))
print('\n-----------------------------------------')
print('TESSERACT OUTPUT --> THRESHOLDED IMAGE')
print('-----------------------------------------')
print(clean_text(mode, pytesseract.image_to_string(thresh, config=custom_config)))
print('\n-----------------------------------------')
print('TESSERACT OUTPUT --> OPENED IMAGE')
print('-----------------------------------------')
print(clean_text(mode, pytesseract.image_to_string(opening, config=custom_config)))
print('\n-----------------------------------------')
print('TESSERACT OUTPUT --> CANNY EDGE IMAGE')
print('-----------------------------------------')
print(clean_text(mode, pytesseract.image_to_string(canny, config=custom_config)))
print('\n-----------------------------------------')
print('TESSERACT OUTPUT --> ZOOMed IMAGE')
print('-----------------------------------------')
print(clean_text(mode, pytesseract.image_to_string(image2, config=custom_config)))