import cv2
from fastai import *
from fastai.vision import *
learn = load_learner('.', 'export.pkl')

# Load the cascade
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

# To capture video from webcam.
cap = cv2.VideoCapture(0)
# To use a video file as input
# cap = cv2.VideoCapture('filename.mp4')
i=0
while True:
    # Read the frame
    _, img = cap.read()
    # Convert to grayscale
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    #print(gray.shape)
    # Detect the faces
    faces = face_cascade.detectMultiScale(gray, 1.1, 5, minSize=(48, 48))
    # Draw the rectangle around each face
    for (x, y, w, h) in faces:
        cv2.rectangle(img, (x, y), (x+w, y+h), (255, 0, 0), 2)
    # Display
    if faces is ():
        cv2.putText(img, "No Face Found", (20, 60) , cv2.FONT_HERSHEY_SIMPLEX,2, (0,255,0), 3)
    else:
        roi = gray[y:y+h, x:x+w]
        roi_resize = cv2.resize(roi, (64, 64))
        # cv2.imwrite('x'+str(i)+'.jpg',roi)
        i+=1
        img_fastai = Image(pil2tensor(roi_resize, dtype=np.float32).div_(255))
        print(learn.predict(img_fastai))


    cv2.imshow('img', img)
    # Stop if escape key is pressed

    k = cv2.waitKey(30) & 0xff
    if k==27:
        break
# Release the VideoCapture object
cap.release()