import cv2
import numpy as np
import keras

'''
задачи
1) находить контур цифры
2) получать его координаты, вывести отдельным изображением
3) масштабировать его под нужный размер
4) загнать это на нейронную сеть https://habr.com/ru/post/466565/ 
5) вывести результат

очень сильно шумит изображение
'''


def find_number(frame):
    # функция выполняет предобработку изображения
    # находит область с цифрой по контурам
    # далее в цикле перебераем найденные контуры
    # если контур найден, то меняем разрешение на 28*28
    # выполняем бинаризацию, делим на 255 и инвертируем вычитая 1
    # flag нужен, чтобы он не выводил единицу, если цифра не найдена
    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)