66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
import tensorflow as tf
|
|
import numpy as np
|
|
import cv2
|
|
import os
|
|
from tensorflow.keras import layers, models
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
# 1. Load dữ liệu
|
|
def load_data(data_dir):
|
|
images = []
|
|
labels = []
|
|
for label, folder in enumerate(os.listdir(data_dir)):
|
|
folder_path = os.path.join(data_dir, folder)
|
|
for file in os.listdir(folder_path):
|
|
img_path = os.path.join(folder_path, file)
|
|
img = cv2.imread(img_path)
|
|
img = cv2.resize(img, (128, 128)) # Resize về kích thước cố định
|
|
images.append(img)
|
|
labels.append(label)
|
|
images = np.array(images) / 255.0 # Chuẩn hóa
|
|
labels = np.array(labels)
|
|
return images, labels
|
|
|
|
# Đường dẫn dữ liệu
|
|
data_dir = '/home/joseph/DetectFace/dataset'
|
|
images, labels = load_data(data_dir)
|
|
|
|
# Chia dữ liệu train/test
|
|
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)
|
|
|
|
# 2. Tạo mô hình phát hiện khuôn mặt
|
|
model = models.Sequential([
|
|
layers.Conv2D(32, (3, 3), activation='relu', input_shape=(128, 128, 3)),
|
|
layers.MaxPooling2D((2, 2)),
|
|
layers.Conv2D(64, (3, 3), activation='relu'),
|
|
layers.MaxPooling2D((2, 2)),
|
|
layers.Conv2D(128, (3, 3), activation='relu'),
|
|
layers.Flatten(),
|
|
layers.Dense(128, activation='relu'),
|
|
layers.Dense(len(set(labels)), activation='softmax') # Số lớp tương ứng số nhãn
|
|
])
|
|
|
|
# Compile mô hình
|
|
model.compile(optimizer='adam',
|
|
loss='sparse_categorical_crossentropy',
|
|
metrics=['accuracy'])
|
|
|
|
# Huấn luyện mô hình
|
|
model.fit(X_train, y_train, epochs=10, validation_data=(X_test, y_test))
|
|
|
|
# 3. Lưu mô hình
|
|
model.save('face_detection_model.h5')
|
|
|
|
# 4. Sử dụng mô hình để dự đoán
|
|
def predict_face(image_path, model_path='face_detection_model.h5'):
|
|
model = tf.keras.models.load_model(model_path)
|
|
img = cv2.imread(image_path)
|
|
img_resized = cv2.resize(img, (128, 128)) / 255.0
|
|
img_resized = np.expand_dims(img_resized, axis=0)
|
|
prediction = model.predict(img_resized)
|
|
return np.argmax(prediction), np.max(prediction)
|
|
|
|
# Test với một ảnh
|
|
image_path = '/home/joseph/DetectFace/test/NGUYEN\ HOANG\ VI_check\ in_at_2024_09_05_11_31_24.png'
|
|
label, confidence = predict_face(image_path)
|
|
print(f"Detected label: {label}, Confidence: {confidence:.2f}") |