结合连通块平均分割以及投影矫正的验证码分割算法

import cv2,os,sys,math
import numpy as np
from pylab import *
%matplotlib inline

def getBinary(path):
    '''
    输入：图片路径名称
    输出：黑底二值化的图片

    '''
    im=cv2.imread(path,0)
    thresh,im=cv2.threshold(255-im,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    return im

def shadow(im,angel):
    '''
    输入：灰度图片，投影角度
    输出：投影大小

    '''
    x=[]
    height , width = im.shape
    for i in xrange(height):
        for j in xrange(width):
            if im[i][j]==255:
                x.append(j-1.0*i/math.tan(math.radians(angel)))
    x=np.array(x)
    return x.max()-x.min()

def shadowTest():
    '''
    测试shadow函数
    '''
    directory='weibo'
    pics=os.listdir(directory)
    for pic in pics[:1]:
        im=getBinary(os.path.join(directory,pic))
        print shadow(im,50)
        figure()
        gray()
        imshow(im)

def getAngle(im):
    '''
    输入：灰度图
    输出：最优的投影角度
    '''
    minShadow=500
    ans=90
    for angle in np.linspace(45,135,91):
        thisShadow=shadow(im,angle)
        if minShadow>thisShadow:
            ans=angle
            minShadow=thisShadow
    return ans

def getAngleTest():
    '''
    测试getAngle函数
    '''
    directory='weibo'
    pics=os.listdir(directory)
    for pic in pics[:5]:
        im=getBinary(os.path.join(directory,pic))
        print getAngle(im)
        gray()
        figure()
        imshow(im)

def affine(im,angle=None):
    '''
    输入：灰度图，仿射变换角度（默认为自适应最优角度）
    输出：旋转后的灰度图
    '''
    height,width=im.shape
    if angle==None:
        angle=getAngle(im)
    pts1=np.float32([[width/2,height/2],[width/2,0],[0,height/2]])
    pts2=np.float32([[width/2,height/2],[width/2+height/2/math.tan(math.radians(angle)),0],[0,height/2]])
    M=cv2.getAffineTransform(pts1,pts2)
    dst=cv2.warpAffine(im,M,(width,height))
    return dst

def affineTest():
    '''
    测试affine函数
    '''
    directory='weibo'
    pics=os.listdir(directory)
    for pic in pics[:50]:
        im=getBinary(os.path.join(directory,pic))
        dst=affine(im)
        gray() 
        figure()
        imshow(np.hstack([im,dst]))
        axis('off')

def splitImg(im,num):
    '''
    输入：灰度图，图中包含的字符个数
    输出：分割后的图片数组
    '''
    height,width=im.shape
    maxPos=0
    minPos=width
    for i in xrange(height):
        for j in xrange(width):
            if im[i][j]==255:
                maxPos=max(maxPos,j)
                minPos=min(minPos,j)
    points=[]
    ans=[]
    for i in xrange(num+1):
        points.append(minPos+(maxPos-minPos)/num*i)
    for i in xrange(num):
        left=points[i]
        right=points[i+1]
        p1=np.zeros((height,left))
        p2=np.ones((height,right-left+1))*255
        p3=np.zeros((height,width-right))
        new_im=np.hstack([p1,p2,p3])
        for i in xrange(height):
            for j in xrange(width):
                if im[i][j]==0:
                    new_im[i][j]=0
        ans.append(new_im)
    return ans

def seedSplit(im):
    '''
    输入：灰度图
    输出：从图中高亮提取出的字符图片数组
    '''
    angle=getAngle(im)
    height,width = im.shape
    vis=np.zeros((height,width))
    vis2=np.zeros((height,width))
    def mark(i,j):
        vis2[i][j]=1
        points=[]
        points.append((i,j))
        for dx in [-1,0,1]:
            for dy in [-1,0,1]:
                if i+dx>=0 and i+dx<height and j+dy>=0 and j+dy<width and im[i+dx][j+dy]==255 and vis2[i+dx][j+dy]==0:
                     points.extend(mark(i+dx,j+dy))
        return points

    def delete(i,j):
        vis[i][j]=1
        ans=1
        for dx in [-1,0,1]:
            for dy in [-1,0,1]:
                if i+dx>=0 and i+dx<height and j+dy>=0 and j+dy<width and im[i+dx][j+dy]==255 and vis[i+dx][j+dy]==0:
                    ans+=delete(i+dx,j+dy)
        return ans

    def getImg(points):
        new_im=np.zeros((height,width))
        for point in points:
            i,j=point         
            new_im[i][j]=255
        return new_im,len(points)

    imgs=[]
    for i in xrange(height):
        for j in xrange(width):
            if vis[i][j]==1:
                continue
            if im[i][j]==255:
                num=delete(i,j)   
                if num<=50:
                    continue
                points=mark(i,j)
                imgs.append(getImg(points))
            vis[i][j]=1

    imgs.sort(lambda left,right:cmp(left[1],right[1]))
    new_imgs=[]
    for i,j in imgs:
        new_imgs.append((affine(i,angle),j))
    imgs=new_imgs
    ans=[]
    if len(imgs)>4:
        imgs=imgs[len(imgs)-4:]
        for i in imgs[len(imgs)-4:]:
            ans.append(i[0])
    elif len(imgs)==4:
        for i in imgs:
            ans.append(i[0])
    elif len(imgs)==3:
        ans.append(imgs[0][0])
        ans.append(imgs[1][0])
        ans.extend(splitImg(imgs[2][0],2))
    elif len(imgs)==2:
        rate=imgs[0][1]*1.0/imgs[1][1]
        if rate>0.6:
            ans.extend(splitImg(imgs[0][0],2))
            ans.extend(splitImg(imgs[1][0],2))
        else:
            ans.append(imgs[0][0])
            ans.extend(splitImg(imgs[1][0],3))
    else:
        ans.extend(splitImg(imgs[0][0],4))
    return ans

def seedSplitTest():
    '''
    测试seedSplit函数
    '''
    directory='weibo'
    pics=os.listdir(directory)
    gray()
    for pic in pics[5:15]:
        im=getBinary(os.path.join(directory,pic))
        figure()
        imshow(im)
        imgs=seedSplit(im)   
        for i in imgs:
            figure()
            imshow(i)

def getTrainPic(im):
    '''
    输入：图片
    输出：依次输出图片中包含的字符，并归一化成统一大小
    '''
    def getCenterx(img):
        height,width=im.shape
        points=[]
        for i in xrange(height):
            for j in xrange(width):
                if img[i][j]==255:
                    points.append(j)
        return np.average(points)

    def getResult(im,row=30,column=30):
        height,width=im.shape
        maxX=0
        maxY=0
        minX=99999
        minY=99999
        for i in xrange(height):
            for j in xrange(width):
                if im[i][j]==255:
                    maxX=max(maxX,j)
                    minX=min(minX,j)
                    maxY=max(maxY,i)
                    minY=min(minY,i)
        im=im[minY:maxY+1].T[minX:maxX+1].T
        im=cv2.copyMakeBorder(im,3,3,3,3,cv2.BORDER_CONSTANT)
        im=cv2.resize(im,(30,30))
        return im

    imgs=seedSplit(im)
    for i in xrange(len(imgs)):
        imgs[i]=(imgs[i],getCenterx(imgs[i]))
    imgs.sort(lambda x,y:cmp(x[1],y[1]))
    ans=[]
    for i,j in imgs:
        ans.append(getResult(i))
    return ans

def getTrainPicTest():
    '''
    测试getTrain函数
    '''
    directory='weibo'
    pics=os.listdir(directory)
    gray()
    for pic in pics[:10]:
        im=getBinary(os.path.join(directory,pic))
        imgs=getTrainPic(im)
        res=np.hstack([imgs[0],imgs[1],imgs[2],imgs[3]])
        figure()
        imshow(im)
        axis('off')
        figure()
        imshow(res)
        axis('off')

getTrainPicTest()