from cv2 import waitKey
import numpy
import pickle
import cv2
import os
import uuid

#图片预处理
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 feature(A):
    midx=int(A.shape[1]/2)+1
    midy=int(A.shape[0]/2)+1
    A1=cv2.mean(A[0:int(A.shape[0]/5),0:A.shape[1]])
    A2=cv2.mean(A[0:midy,0:int(A.shape[1]/3)])
    A3=cv2.mean(A[0:midy,int(A.shape[1]*2/3):A.shape[1]])
    A4=cv2.mean(A[int(A.shape[0]*3/5):int(A.shape[0]*4/5),0:A.shape[1]])
    A5=cv2.mean(A[midy:A.shape[0],0:int(A.shape[1]/3)])
    A6=cv2.mean(A[midy:A.shape[0],int(A.shape[1]*2/3):A.shape[1]])
    A7=cv2.mean(A[int(A.shape[0]*4/5):A.shape[0],0:A.shape[1]])
    if A1[0] >128 :
        A1=1
    else:
        A1=0
    if A2[0] >128 :
        A2=1
    else:
        A2=0
    if A3[0] >128 :
        A3=1
    else:
        A3=0
    if A4[0] >128 :
        A4=1
    else:
        A4=0
    if A5[0] >128 :
        A5=1
    else:
        A5=0
    if A6[0] >128 :
        A6=1
    else:
        A6=0
    if A7[0] >128 :
        A7=1
    else:
        A7=0
           
    AF=[A1,A2,A3,A4,A5,A6,A7]
    #print(AF)
    i=0
    tube=0
    onenumber=0
    while(i < 7):
        if(AF[i]==1):
            tube = tube + pow(2,i)
        i +=1
    if(tube==63):
        onenumber = 0
    elif(tube==6):
        onenumber = 1
    elif(tube==91):
        onenumber = 2
    elif(tube==79):
        onenumber = 3
    elif(tube==102 or tube==110):
    #110是因为有干扰情况
        onenumber = 4
    elif(tube==109):
        onenumber = 5
    elif(tube==125):
        onenumber = 6
    elif(tube==7):
        onenumber = 7
    elif(tube==127):
        onenumber = 8
    elif(tube==103):
        onenumber = 9
    else:
        onenumber = -1 
               
    return AF,onenumber
 
      

#训练已知图片的特征         
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('./data/'+str(i)+'.jpg',0)
            AF,nmber = 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):
    key=knn(train_set,AF,k)
    return key

#识别文件夹中的所有图片上的数据并输出到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中
    path = "./images4"
    i=0
    filename = os.listdir(path)    # 获取文件夹文件
    for file in filename:
        img = cv2.imread(path + "/" + file,0)    # 读取图片
        AF,nmber=feature(img)
        print(AF,nmber)
        result=identification(train_set,AF,k)
        cv2.imwrite('./images4_infer/'+str(result)+'_'+str(uuid.uuid4())+".jpg",img)
        i = i+ 1
        if i >200 :
            break

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