import time import cv2 import numpy as np from matplotlib import pyplo

import time
import cv2
import numpy as np
from matplotlib import pyplot as plt
from collections import Counter


import matplotlib
matplotlib.rcParams['savefig.dpi'] = 4 * matplotlib.rcParams['savefig.dpi']

img = cv2.imread('../tessdata/1.png', cv2.CV_LOAD_IMAGE_GRAYSCALE)[245:727+245, 90:1198+90]

import math
def q(x):
    return math.sqrt(x*x*x)

def findTextColor(img):
    tmp = np.copy(img[100:200,100:200])
    tmp[tmp < 100] = 0

    contours, _ = cv2.findContours(tmp, cv2.RETR_LIST, cv2.RETR_LIST )

    colors = []
    for c in contours:
        region = tmp[c]
        region = region[region > 100]
        colors.append(max(region.flatten()))
        #colors.append(np.sum(region) / len(region))
    return max(colors)
#    return np.sum(colors) / len(colors)

print 'text color:',findTextColor(img)

def invertSelectionLine(img):
    def detectSelectionLine(img):
        for i in range(img.shape[0]):
            if np.sum(img[i,:]) / img.shape[1] > 100:
                for j in range(i, img.shape[0]):
                    if np.sum(img[j,:]) / img.shape[0] <= 100:
                        return i, j
                return i, img.shape[0]
    firstRow, lastRow = detectSelectionLine(img)
    textColor = findTextColor(img)
    inv = img[firstRow:lastRow,:].astype('int16')
    inv = 255 - inv

    globalAvg = np.average(img)
    areaAvg = np.average(inv)
    inv -= areaAvg - globalAvg
    inv[inv<0]=0
    img[firstRow:lastRow,:] = inv.astype('uint8')

def normalizeColors(img):
    tmp = (img.astype('int16') * 1.2)
    maxColor = max(tmp.flatten())
    maxColor -= 255
    tmp -= maxColor
    tmp[tmp<0] = 0
    tmp *= 3
    tmp[tmp>255] = 255
    return img


def contrast(img):
    THRESHOLD = 80
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if img[i,j] > THRESHOLD:
                img[i,j] += min(255 - img[i,j], q(img[i,j]-THRESHOLD))
            else:
                img[i,j] -= min(img[i,j], q(THRESHOLD - img[i,j]))

def contrast(img, mul=2, shift=0):
    tmp = np.copy(img).astype(float) * mul
    avg = np.sum(tmp.flatten()) / len(tmp.flatten()) + shift
    tmp -= avg
    tmp[tmp < 0] = 0
    tmp[tmp > 255] = 255
    return tmp.astype('uint8')

def detectVerticalLines(img):
    for col in range(2, img.shape[1]):
        diff = np.sum(abs(img[:,col] - img[:,col-1])) / img.shape[0]
        if diff > 200:
            for col2 in range(col+1, img.shape[1]):
                diff2 = np.sum(abs(img[:,col2] - img[:,col2-1])) / img.shape[0]
                if diff2 > 200:
                    width = col2 - col
                    hw = width / 2
                    if hw == 0: hw = 1

                    mid = (col2 + col) / 2

                    img[:,col:mid] = img[:,col-hw-1:col-1]
                    img[:,(col+col2)/2:col2] = img[:,col2+1:col2+1+hw]
                    return

#for i in range(3):
#    detectVerticalLines(img)

def removeVerticalLines(img):
    minrows = [r[1] for r in sorted([(np.sum(img[row,:]),row) for row in range(img.shape[0])])[:50]]
    cols = []
    for col in range(100,img.shape[1]):
        avg = np.sum(img[minrows,col]) / len(minrows)
        cols.append(avg)

    freq = np.fft.fft(cols)
    freq[100:] = 0
    diff = abs(cols - np.fft.ifft(freq))
    maxDiff = max(diff)
    for i in range(len(diff)):
        if diff[i] > maxDiff *0.20:
            img[:,i+100] = 0#np.sum(img.flatten()) / len(img.flatten())

def removeHorizontalLines(img):
    mincols = [r[1] for r in sorted([(np.sum(img[:,col]),col) for col in range(img.shape[1])])[:50]]
    rows = []
    for row in range(100,img.shape[0]):
        avg = np.sum(img[row, mincols]) / len(mincols)
        rows.append(avg)

    freq = np.fft.fft(rows)
    freq[100:] = 0
    diff = abs(rows - np.fft.ifft(freq))
    maxDiff = max(diff)
    for i in range(len(diff)):
        if diff[i] > maxDiff *0.20:
            img[i+100,:] = 0#np.sum(img.flatten()) / len(img.flatten())




#img = contrast(img, 2, 40)
#img = contrast(img, 1.5, 20)
#img = contrast(img, 1.5, 10)
#img[img<10] = 0
#img[img>0] = 255

def fftRows(img):
    for i in range(img.shape[0]):
        row = img[i]
        dark = row[row < 100]
        fimg = np.fft.fft(dark)
        fimg[0] = 0
        row[row < 100] = abs(np.fft.ifft(fimg))


def cleanBackgroundFFT(img):
    print 'fft start'
    start = time.time()
    dark = img[img < 100]
    chunkSize = img.shape[1] * 10
    chunkCount = int(math.ceil(float(len(dark))/chunkSize)*chunkSize)
    for chunk in range(0, chunkCount, chunkSize):
        fimg = np.fft.fft(dark[chunk:chunk+chunkSize])
        fimg[:10] = 0
        dark[chunk:chunk+chunkSize] = abs(np.fft.ifft(fimg))
    img[img<100] = dark
    end = time.time()
    print 'fft end'
    print 'fft took', end-start


#img[row,:] = abs(np.fft.ifft(fimg))

#fimg = np.fft.fft2(img)
#fimg[0] /= 1000
#fimg[1] /= 1000
#img = abs(np.fft.ifft2(fimg))


invertSelectionLine(img)
removeHorizontalLines(img)
cleanBackgroundFFT(img)
removeVerticalLines(img)
img = contrast(img, 3, 20)
img = contrast(img, 1.5, 10)

cv2.imwrite('q.tiff', img)
#plt.imshow(img, 'gray')