import numpy
import pickle
import cv2

#图片预处理
def pretreatment(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    thresh = cv2.threshold(gray, 0, 255,cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
    return thresh

#提取图片特征
def feature(A):
    midx=int(A.shape[1]/2)+1
    midy=int(A.shape[0]/2)+1
    A1=cv2.mean(A[0:midy,0:midx])
    A2=cv2.mean(A[midy:A.shape[0],0:midx])
    A3=cv2.mean(A[0:midy,midx:A.shape[1]])
    A4=cv2.mean(A[midy:A.shape[0],midx:A.shape[1]])
    A5=cv2.mean(A)
    AF=[A1,A2,A3,A4,A5]
    return AF


#训练已知图片的特征         
def training():
    train_set={}
    for i in range(11):
        value=[]
        for j in range(10):
            print('./'+str(i)+'/'+str(i)+'-'+str(j)+'.jpg')
            ima=cv2.imread('./'+str(i)+'/'+str(i)+'-'+str(j)+'.jpg',0)
            if ima is None:
                print('image is none !')
            print(ima.shape)
            AF = feature(ima)
            value.append(AF)
            #print(value)
        train_set[i]=value
    #把训练结果存为永久文件，以备下次使用
    output=open('./train_set.pkl','wb')
    pickle.dump(train_set,output)
    output.close()
    return train_set
    
#计算两向量的距离
def distance(v1,v2):
    vector1=numpy.array(v1)
    vector2=numpy.array(v2) 
    Vector=(vector1-vector2)**2
    distance = Vector.sum()**0.5
    return distance

#用最近邻算法识别单个数字
def knn(train_set,V,k):
    key_sort=[11]*k
    value_sort=[11]*k
    for key in range(11):
        for value in train_set[key]:
            d=distance(V,value)
            for i in range(k):#从小到大排序
                if d<value_sort[i]:
                    for j in range(k-2,i-1,-1):
                        key_sort[j+1]=key_sort[j]
                        value_sort[j+1]=value_sort[j]
                    key_sort[i]=key
                    value_sort[i]=d
                    break
    max_key_count=-1
    key_set=set(key_sort)
    for key in key_set:   #统计每个类别出现的次数
        if max_key_count<key_sort.count(key):
            max_key_count=key_sort.count(key)
            max_key=key
    return max_key

#合并一副图片的所有数字
def identification(train_set,AF,k):
    result=''
    for i in AF:
        key=knn(train_set,i,k)
        if key==10:
            key='.'
        result=result+str(key)
    return float(result)

#识别文件夹中的所有图片上的数据并输出到EXCELL中
def main(k,trained=None):
    # k：knn算法的的k，即取训练集中最相邻前k个样本进行判别
    #trained：trained=0则程序会开始训练出训练集，否则会用已保存的训练集
    if trained==0:
        train_set=training()
    else:
        pkl_file=open('train_set.pkl','rb')
        train_set=pickle.load(pkl_file)
        pkl_file.close()
    #识别所有图片的数字并输出到EXCELL中
    for i in range(1,21):
        ima=cv2.imread('./test/'+str(i)+'.jpg')
        #im=pretreatment(ima)
        AF=feature(ima)
        result=identification(train_set,AF,k)

main(4,0)  #表示：4近邻，还没有训练集备份，还没有昨日数据备份