人工智能-人脸识别

人脸识别

import cv2, os, numpy as np

# 共享配置

recognizer = cv2.face.LBPHFaceRecognizer_create()

dataset_dir = "dataset"

face_cascade = cv2.CascadeClassifier(os.getcwd() + '\\' + 'haarcascade_frontalface_default.xml')

def collect_data():

    user_id = input("用户ID(数字): ")

    user_dir = f"{dataset_dir}/User_{user_id}_{input('用户名: ')}"

    os.makedirs(user_dir, exist_ok=True)

    cap = cv2.VideoCapture(0,cv2.CAP_DSHOW)

   

    count = 0

    while count < 50 and cv2.waitKey(1) != 27:

        _, frame = cap.read()

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        faces = face_cascade.detectMultiScale(gray, 1.3, 5)

        for x, y, w, h in faces:

            cv2.imwrite(

                f"{user_dir}/{count}.jpg",

                cv2.resize(gray[y : y + h, x : x + w], (200, 200)),

            )

            cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)

            count += 1

        cv2.imshow("Collecting...", frame)

    cap.release()

    cv2.destroyAllWindows()

 

def train_model():

    faces, ids = [], []

    for root, _, files in os.walk(dataset_dir):

        for file in files:

            if file.endswith(".jpg"):

                path = os.path.join(root, file)

                faces.append(cv2.imread(path, 0))

                ids.append(int(os.path.basename(root).split("_")[1]))

    recognizer.train(faces, np.array(ids))

    recognizer.save("model.yml")

 

def recognize():

    id_map = {0: "Unknown"}

    for d in os.listdir(dataset_dir):

        if d.startswith("User_"):

            id_map[int(d.split("_")[1])] = d.split("_")[2]

    recognizer.read("model.yml")

    cap = cv2.VideoCapture(0,cv2.CAP_DSHOW)

    while cv2.waitKey(1) != 27:

        _, frame = cap.read()

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        for x, y, w, h in face_cascade.detectMultiScale(gray, 1.3, 5):

            label, conf = recognizer.predict(

                cv2.resize(gray[y : y + h, x : x + w], (200, 200))

            )

            name = id_map.get(label, "Unknown") if conf < 60 else "Unknown"

            color = (0, 255, 0) if name != "Unknown" else (0, 0, 255)

            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)

            cv2.putText(frame, name, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)

        cv2.imshow("Recognizing...", frame)

    cap.release()

    cv2.destroyAllWindows()

 

if __name__ == "__main__":

    func = [collect_data, train_model, recognize][

        int(input("功能选择 (1采集 2训练 3识别): ")) - 1

    ]

    func()