import cv2
import numpy as np
import keras

def find_number(frame):
    # функция выполняет предобработку изображения
    # находит область с цифрой по контурам
    # далее в цикле перебераем найденные контуры применяя фильтры
    # если контур найден, то меняем разрешение на 28*28
    # выполняем бинаризацию, делим на 255 и инвертируем вычитая 1
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray,(3,3), 0)
    edges = cv2.Canny(blur, 50, 100)
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    img2 = np.zeros((28, 28), np.uint8)
    img_rec = np.zeros((1, 28, 28, 1), np.uint8)
    fail = True

    for contr in contours:

        if cv2.contourArea(contr) < 500:
            continue
        x,y,w,h = cv2.boundingRect(contr)
        img2 = frame[y-5:y+h+5, x-5:x+w+5]

        # Проверка на некорректное изображение
        if img2.shape[0] <= 0 or img2.shape[1]<= 0:
            continue

        #Фильтер по вертикальности рамки
        if img2.shape[0] < img2.shape[1]:
            continue

        img2 = cv2.resize(img2, (28, 28))
        img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

        # th3 = cv2.adaptiveThreshold(img2,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\
        #     cv2.THRESH_BINARY,11,2)

        _, th3 = cv2.threshold(img2,127,255,cv2.THRESH_BINARY)

        # Фильтруем по следнему цвету, должно быть много белого
        avg_color_per_row = np.average(th3, axis=0)
        avg_color = np.average(avg_color_per_row, axis=0)
        if avg_color < 180:
            continue
        
        # Если все норм то рисуем рамки на исходном изображении.
        img = cv2.rectangle(frame,(x,y),(x+w,y+h),(0,255,0),2)

        cv2.imshow("img2", img2)
        cv2.imshow("th3", th3)

        img_rec = th3/255.0
        img_rec = 1 - img_rec
        img_rec = img_rec.reshape((1,28,28,1))

        fail = False
    return img2, img_rec, fail



if __name__ == '__main__':
    cam = cv2.VideoCapture(0)
    model = keras.models.load_model("mnist_trained_model.h5")
    try:
        while True:
            _, frame = cam.read()
            img_show, img_rec, fail = find_number(frame)
            if fail == False:
                result = model.predict_classes([img_rec])
                cv2.putText(frame,str(result[0]),(10,460), cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),2,cv2.LINE_AA)
                

            cv2.imshow("Image", frame)
            
            k = cv2.waitKey(2)
            if k == 27:
                #cv2.imwrite('diff_robot_nn/nine.jpg', img_show)
                cv2.destroyAllWindows()
                break
    except Exception as e:
        cam.release()
        cv2.destroyAllWindows()
        print(e)