from opencv.cv import *
from opencv.highgui import *
import sys

SCALE = 2
WIDTH = 1024/SCALE
HEIGHT = 768/SCALE
THRESH_PARAM = 20

cvNamedWindow("ocr")

originalImg = cvLoadImage(sys.argv[1], 0)

scaledImg = cvCreateImage(cvSize(WIDTH, HEIGHT), 8, 1)
cvResize(originalImg, scaledImg)


binImg = cvClone(scaledImg)
#RADIUS = 20;
#elt = cvCreateStructuringElementEx(2*RADIUS+1, 2*RADIUS+1, RADIUS, RADIUS, CV_SHAPE_ELLIPSE, None)
#tmpImg = cvCreateImage(cvSize(WIDTH, HEIGHT), 8, 1)
#cvMorphologyEx(binImg, binImg, tmpImg, elt, CV_MOP_BLACKHAT)
cvAdaptiveThreshold(binImg, binImg, 255, 0, CV_THRESH_BINARY_INV, 41, THRESH_PARAM)

   
storage = cvCreateMemStorage(0)
lines = cvHoughLines2(binImg, storage, CV_HOUGH_STANDARD, 1, CV_PI/180, 100)
for i in xrange(min(lines.total,10)):
    print "theta = ", lines[i][1]*180/CV_PI
theta = lines[0][1]*180/CV_PI

rotImg = cvCreateImage(cvSize(WIDTH, HEIGHT), 8, 1)
tmpImg = cvCreateImage(cvSize(WIDTH, HEIGHT), 8, 1)
map = cvCreateMat(2, 3, CV_32FC1)
cv2DRotationMatrix(cvPoint(WIDTH/2, HEIGHT/2), 90-theta, 1, map)
cvAdaptiveThreshold(scaledImg, tmpImg, 255, 0, CV_THRESH_BINARY, 41, THRESH_PARAM)
cvWarpAffine(tmpImg, rotImg, map, CV_WARP_FILL_OUTLIERS+CV_WARP_INVERSE_MAP, 255)
cvAdaptiveThreshold(rotImg, rotImg, 255, 0, CV_THRESH_BINARY, 41, THRESH_PARAM)


blobImg = cvCloneImage(rotImg)
ELLIPSE_RADIUS = 2;
RECT_RADIUS = 5/SCALE;
eltRect = cvCreateStructuringElementEx(2*RECT_RADIUS+1, 2*RECT_RADIUS+1, RECT_RADIUS, RECT_RADIUS, CV_SHAPE_RECT, None)
#cvMorphologyEx(blobImg, blobImg, tmpImg, eltRect, CV_MOP_OPEN, 2)
#cvMorphologyEx(blobImg, blobImg, tmpImg, eltRect, CV_MOP_CLOSE, 1)
cvErode(blobImg, blobImg, eltRect, 1)

def extendRect (blobImg, r):
    try:
        xtop, ytop, xbot, ybot = r
        for x in range(xtop,xbot+1):
            if blobImg[ytop-1,x] < 128:
                ytop -= 1
                return True, (xtop,ytop,xbot,ybot)
            if blobImg[ybot+1,x] < 128:
                ybot += 1
                return True, (xtop,ytop,xbot,ybot)
        for y in range(ytop,ybot+1):
            if blobImg[y,xbot+1] < 128:
                xbot += 1
                return True, (xtop,ytop,xbot,ybot)
            if blobImg[y,xtop-1] < 128:
                xtop -= 1
                return True, (xtop,ytop,xbot,ybot)
    except IndexError:
        pass
    return False, r

def findRect (blobImg, x0, y0):
    # Find extremal points
    r = (x0-2, y0-2, x0+2, y0+2)
    while True:
        extended, r = extendRect(blobImg, r)
        if not extended:
            break
    xtop, ytop, xbot, ybot = r
    print xtop,ytop,"-",xbot,ybot
    pad = 10/SCALE
    return (cvPoint(xtop-pad,ytop-pad), cvPoint(xbot+pad,ybot+pad))

top, bot = findRect(blobImg, WIDTH/2, HEIGHT/2)
cvRectangle(rotImg, top, bot, CV_RGB(0,255,0), 2)
cvRectangle(blobImg, top, bot, CV_RGB(0,255,0), 2)

cvShowImage("ocr", binImg)
cvWaitKey()
cvShowImage("ocr", rotImg)
cvWaitKey()
cvShowImage("ocr", blobImg)
cvWaitKey()