import cv2

ESC = 27			# 設定將以 ESC 鍵作為離開程式		
n = 1				# 畫面數量計數
index = 0			# 存檔檔名用	
total = 100			# 人臉取樣總數

# 自訂函數，images 與 h0 要自己先建立。若有第二個人就建立 h1
def saveImage(face_image, index):
    filename = 'images/h0/{:03d}.pgm'.format(index)
    cv2.imwrite(filename, face_image)
    print(filename)

# 載入聯集分類器與開啟攝影機
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
cap = cv2.VideoCapture(0)
ratio = cap.get(cv2.CAP_PROP_FRAME_WIDTH) / cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
WIDTH = 400
HEIGHT = int(WIDTH / ratio)
cv2.namedWindow('video', cv2.WINDOW_NORMAL)

# 讀取影像並轉成灰階
while n > 0:
    ret, frame = cap.read()
    frame = cv2.resize(frame, (WIDTH, HEIGHT))
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)		# 轉成灰階
    
    # 偵測人臉，並且每 5 張人臉存檔一次
    faces = face_cascade.detectMultiScale(gray, 1.1, 3)	# 辨識人臉

    for (x, y, w, h) in faces:
        frame = cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 3)

        if n % 5 == 0:							# 每 5 張存檔一次，目的是要讓使用者變換姿勢
            face_img = gray[y: y + h, x: x + w]
            face_img = cv2.resize(face_img, (400, 400))	# 訓練用圖片解析度為 400x400，最小為 50x50
            saveImage(face_img, index)				# 儲存訓練圖片
            index += 1

            if index >= total:					# 若拍照滿 100 張，則跳離 for 迴圈
                print('get training data done')
                n = -1
                break

        n += 1
    
    # 將攝影機拍到的影像顯示出來
    cv2.imshow('video', frame)

    if cv2.waitKey(1) == ESC:
        cv2.destroyAllWindows()
        break