321 lines
12 KiB
Python
321 lines
12 KiB
Python
import os
|
||
import sys
|
||
import argparse
|
||
import numpy as np
|
||
import json
|
||
import glob
|
||
import time
|
||
import cv2
|
||
from mtcnn.mtcnn import MTCNN
|
||
import kp
|
||
|
||
# 引入之前的功能
|
||
from nef_test import (
|
||
load_image_safe, landmarks, affine_matrix, extract_vector_data,
|
||
load_vector, cosine_similarity, SCPU_FW_PATH, NCPU_FW_PATH, visualize_alignment
|
||
)
|
||
|
||
# 預設參數
|
||
MODEL_FILE_PATH = 'R34_G369K.nef'
|
||
VECTOR_DATABASE_DIR = 'face_vectors'
|
||
SIMILARITY_THRESHOLD = 0.5 # 相似度閾值，可根據需要調整
|
||
|
||
def get_face_vector(device_group, model_nef_descriptor, image_path):
|
||
"""從圖像中擷取人臉向量"""
|
||
|
||
# 創建MTCNN檢測器
|
||
detector = MTCNN(device="CPU:0")
|
||
|
||
# 載入圖像
|
||
try:
|
||
img_rgb, img_bgr = load_image_safe(image_path)
|
||
print(f" - 已載入圖像: {image_path}")
|
||
except Exception as e:
|
||
print(f"錯誤: {str(e)}")
|
||
return None
|
||
|
||
# 獲取人臉特徵點並計算仿射變換矩陣
|
||
try:
|
||
lmks = landmarks(detector, img_rgb)
|
||
mat, size = affine_matrix(lmks)
|
||
print(" - 已檢測到人臉特徵點")
|
||
except Exception as e:
|
||
print(f"錯誤: {str(e)}")
|
||
return None
|
||
|
||
# 應用仿射變換
|
||
aligned_img = cv2.warpAffine(img_bgr, mat, size)
|
||
|
||
# visualize_alignment(img_bgr, lmks, aligned_img)
|
||
|
||
# 轉換為BGR565格式並調整大小
|
||
aligned_img_bgr565 = cv2.cvtColor(aligned_img, cv2.COLOR_BGR2BGR565)
|
||
img_bgr565 = cv2.resize(aligned_img_bgr565, (112, 112), interpolation=cv2.INTER_LINEAR)
|
||
print(" - 已對齊圖像並格式化")
|
||
|
||
# 準備通用圖像推理輸入描述符
|
||
generic_inference_input_descriptor = kp.GenericImageInferenceDescriptor(
|
||
model_id=model_nef_descriptor.models[0].id,
|
||
inference_number=0,
|
||
input_node_image_list=[
|
||
kp.GenericInputNodeImage(
|
||
image=img_bgr565,
|
||
image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
|
||
resize_mode=kp.ResizeMode.KP_RESIZE_ENABLE,
|
||
padding_mode=kp.PaddingMode.KP_PADDING_CORNER,
|
||
normalize_mode=kp.NormalizeMode.KP_NORMALIZE_KNERON
|
||
)
|
||
]
|
||
)
|
||
|
||
# 開始推理工作
|
||
try:
|
||
kp.inference.generic_image_inference_send(
|
||
device_group=device_group,
|
||
generic_inference_input_descriptor=generic_inference_input_descriptor
|
||
)
|
||
|
||
generic_raw_result = kp.inference.generic_image_inference_receive(device_group=device_group)
|
||
print(" - 推理成功完成")
|
||
except kp.ApiKPException as exception:
|
||
print(f' - 錯誤: 推理失敗，錯誤 = {exception}')
|
||
return None
|
||
|
||
# 獲取推理節點輸出
|
||
inf_node_output_list = []
|
||
for node_idx in range(generic_raw_result.header.num_output_node):
|
||
inference_float_node_output = kp.inference.generic_inference_retrieve_float_node(
|
||
node_idx=node_idx,
|
||
generic_raw_result=generic_raw_result,
|
||
channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW
|
||
)
|
||
inf_node_output_list.append(inference_float_node_output)
|
||
|
||
# 獲取人臉向量
|
||
if len(inf_node_output_list) > 0:
|
||
face_vector = inf_node_output_list[0]
|
||
# 轉換向量為標準NumPy數組
|
||
face_vector_np = extract_vector_data(face_vector)
|
||
|
||
# Ensure the vector is flattened to 1D
|
||
face_vector_np = face_vector_np.flatten()
|
||
|
||
# Print shape for debugging
|
||
print(f" - Face vector shape: {face_vector_np.shape}")
|
||
|
||
return face_vector_np
|
||
else:
|
||
print("錯誤: 推理結果中沒有找到輸出節點")
|
||
return None
|
||
|
||
def load_face_database(database_dir):
|
||
"""載入人臉資料庫"""
|
||
face_db = []
|
||
|
||
# 確保資料庫目錄存在
|
||
if not os.path.exists(database_dir):
|
||
print(f"警告: 人臉資料庫目錄 '{database_dir}' 不存在，將創建該目錄")
|
||
os.makedirs(database_dir)
|
||
return face_db
|
||
|
||
# 查找所有支持的向量文件
|
||
vector_files = []
|
||
vector_files.extend(glob.glob(os.path.join(database_dir, "*.npy")))
|
||
vector_files.extend(glob.glob(os.path.join(database_dir, "*.json")))
|
||
vector_files.extend(glob.glob(os.path.join(database_dir, "*.pkl")))
|
||
|
||
# 載入每個向量文件
|
||
for vector_file in vector_files:
|
||
try:
|
||
# 根據文件擴展名自動確定格式
|
||
file_ext = os.path.splitext(vector_file)[1].lower()
|
||
if file_ext == '.npy':
|
||
vector = load_vector(vector_file, format='numpy')
|
||
# 嘗試從文件名中提取標識信息
|
||
name = os.path.basename(vector_file).replace('.npy', '')
|
||
metadata = {'name': name}
|
||
face_db.append({
|
||
'vector': vector,
|
||
'metadata': metadata,
|
||
'file_path': vector_file
|
||
})
|
||
elif file_ext == '.json':
|
||
vector, metadata = load_vector(vector_file, format='json')
|
||
face_db.append({
|
||
'vector': vector,
|
||
'metadata': metadata or {'name': os.path.basename(vector_file).replace('.json', '')},
|
||
'file_path': vector_file
|
||
})
|
||
elif file_ext == '.pkl':
|
||
vector = load_vector(vector_file, format='pickle')
|
||
# 嘗試從文件名中提取標識信息
|
||
name = os.path.basename(vector_file).replace('.pkl', '')
|
||
metadata = {'name': name}
|
||
face_db.append({
|
||
'vector': vector,
|
||
'metadata': metadata,
|
||
'file_path': vector_file
|
||
})
|
||
|
||
print(f"已載入人臉向量: {vector_file}")
|
||
except Exception as e:
|
||
print(f"警告: 無法載入向量文件 '{vector_file}': {str(e)}")
|
||
|
||
print(f"成功載入 {len(face_db)} 個人臉向量")
|
||
return face_db
|
||
|
||
def recognize_face(new_face_vector, face_database, threshold=SIMILARITY_THRESHOLD):
|
||
"""識別人臉，將新的人臉向量與資料庫中的向量進行比較"""
|
||
if not face_database:
|
||
return None, 0.0
|
||
|
||
max_similarity = 0.0
|
||
best_match = None
|
||
|
||
for face_entry in face_database:
|
||
stored_vector = face_entry['vector']
|
||
similarity = cosine_similarity(new_face_vector, stored_vector)
|
||
|
||
if similarity > max_similarity:
|
||
max_similarity = similarity
|
||
best_match = face_entry
|
||
|
||
# 如果最高相似度超過閾值，則返回匹配結果
|
||
if max_similarity >= threshold:
|
||
return best_match, max_similarity
|
||
else:
|
||
return None, max_similarity
|
||
|
||
def add_face_to_database(face_vector, database_dir, name=None, image_path=None, format='json'):
|
||
"""添加新的人臉向量到資料庫"""
|
||
if not os.path.exists(database_dir):
|
||
os.makedirs(database_dir)
|
||
|
||
# 創建唯一的文件名
|
||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
if name:
|
||
file_name = f"{name}_{timestamp}"
|
||
else:
|
||
file_name = f"unknown_{timestamp}"
|
||
|
||
# 創建元數據
|
||
metadata = {
|
||
'name': name or 'unknown',
|
||
'timestamp': time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
'image_path': image_path
|
||
}
|
||
|
||
# 根據格式保存向量
|
||
if format == 'numpy':
|
||
file_path = os.path.join(database_dir, f"{file_name}.npy")
|
||
np.save(file_path, face_vector)
|
||
elif format == 'pickle':
|
||
import pickle
|
||
file_path = os.path.join(database_dir, f"{file_name}.pkl")
|
||
with open(file_path, 'wb') as f:
|
||
pickle.dump(face_vector, f)
|
||
else: # 預設使用JSON格式
|
||
file_path = os.path.join(database_dir, f"{file_name}.json")
|
||
data = {
|
||
'vector': face_vector.tolist(),
|
||
'metadata': metadata
|
||
}
|
||
with open(file_path, 'w') as f:
|
||
json.dump(data, f)
|
||
|
||
print(f"已將新的人臉向量保存到: {file_path}")
|
||
return file_path
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='人臉識別和比對系統')
|
||
parser.add_argument('-p', '--port_id', help='使用指定的端口ID連接設備', default=28, type=int)
|
||
parser.add_argument('-m', '--model', help=f'模型文件路徑 (.nef) (預設: {MODEL_FILE_PATH})',
|
||
default=MODEL_FILE_PATH, type=str)
|
||
parser.add_argument('-i', '--img', help='待識別的圖像文件路徑', required=True, type=str)
|
||
parser.add_argument('-d', '--database', help=f'人臉向量資料庫目錄 (預設: {VECTOR_DATABASE_DIR})',
|
||
default=VECTOR_DATABASE_DIR, type=str)
|
||
parser.add_argument('-t', '--threshold', help=f'相似度閾值 (預設: {SIMILARITY_THRESHOLD})',
|
||
default=SIMILARITY_THRESHOLD, type=float)
|
||
parser.add_argument('-a', '--add', help='將新人臉添加到資料庫 (如果未識別)', action='store_true')
|
||
parser.add_argument('-n', '--name', help='新人臉的名稱 (與 --add 一起使用)', default=None, type=str)
|
||
parser.add_argument('-f', '--format', help='向量儲存格式: numpy, pickle, 或 json (預設: json)',
|
||
default='json', choices=['numpy', 'pickle', 'json'], type=str)
|
||
args = parser.parse_args()
|
||
|
||
# 連接設備
|
||
try:
|
||
print('[連接設備]')
|
||
device_group = kp.core.connect_devices(usb_port_ids=[args.port_id])
|
||
print(' - 成功')
|
||
except kp.ApiKPException as exception:
|
||
print(f'錯誤: 連接設備失敗, 端口ID = \'{args.port_id}\'，錯誤信息: [{str(exception)}]')
|
||
return
|
||
|
||
# 設置USB通信超時
|
||
print('[設置設備超時]')
|
||
kp.core.set_timeout(device_group=device_group, milliseconds=5000)
|
||
print(' - 成功')
|
||
|
||
# 上傳固件到設備
|
||
try:
|
||
print('[上傳固件]')
|
||
kp.core.load_firmware_from_file(device_group=device_group,
|
||
scpu_fw_path=SCPU_FW_PATH,
|
||
ncpu_fw_path=NCPU_FW_PATH)
|
||
print(' - 成功')
|
||
except kp.ApiKPException as exception:
|
||
print(f'錯誤: 上傳固件失敗，錯誤 = \'{str(exception)}\'')
|
||
return
|
||
|
||
# 上傳模型到設備
|
||
try:
|
||
print('[上傳模型]')
|
||
model_nef_descriptor = kp.core.load_model_from_file(device_group=device_group,
|
||
file_path=args.model)
|
||
print(' - 成功')
|
||
except kp.ApiKPException as exception:
|
||
print(f'錯誤: 上傳模型失敗，錯誤 = \'{str(exception)}\'')
|
||
return
|
||
|
||
# 載入人臉資料庫
|
||
print('[載入人臉資料庫]')
|
||
face_database = load_face_database(args.database)
|
||
|
||
# 從圖像中提取人臉向量
|
||
print('[處理輸入圖像]')
|
||
new_face_vector = get_face_vector(device_group, model_nef_descriptor, args.img)
|
||
|
||
if new_face_vector is not None:
|
||
# 識別人臉
|
||
print('[識別人臉]')
|
||
match, similarity = recognize_face(new_face_vector, face_database, args.threshold)
|
||
|
||
if match:
|
||
metadata = match['metadata']
|
||
name = metadata.get('name', '未知')
|
||
print(f"[結果] 識別為: {name}")
|
||
print(f" - 相似度: {similarity:.4f}")
|
||
print(f" - 來源文件: {match['file_path']}")
|
||
if 'timestamp' in metadata:
|
||
print(f" - 記錄時間: {metadata['timestamp']}")
|
||
else:
|
||
print(f"[結果] 未能識別人臉 (最高相似度: {similarity:.4f}, 閾值: {args.threshold})")
|
||
|
||
# 如果指定了--add參數，則將新的人臉添加到資料庫
|
||
if args.add:
|
||
print('[添加新人臉到資料庫]')
|
||
file_path = add_face_to_database(
|
||
new_face_vector,
|
||
args.database,
|
||
name=args.name,
|
||
image_path=args.img,
|
||
format=args.format
|
||
)
|
||
print(f"[添加完成] 已添加新人臉: {args.name or '未命名'}")
|
||
|
||
# 清理
|
||
kp.core.disconnect_devices(device_group=device_group)
|
||
print("[清理] 設備已斷開連接")
|
||
|
||
if __name__ == '__main__':
|
||
main() |