recognition.py

import argparse
import pickle
import cv2
import torch
import numpy as np
import warnings
from insightface.app import FaceAnalysis
from train_arcface import SimpleClassifier

warnings.filterwarnings('ignore')

parser = argparse.ArgumentParser()
parser.add_argument('--image', type=str, help='Path to image for single recognition')
parser.add_argument('--threshold', type=float, default=0.7, help='Confidence threshold')
args = parser.parse_args()

device = 'cuda' if torch.cuda.is_available() else 'cpu'
with open('label_to_id.pkl','rb') as f: label2id = pickle.load(f)
id2label = {v:k for k,v in label2id.items()}

# 모델 로드
classifier = SimpleClassifier(512, len(label2id)).to(device)
ckpt = torch.load('classifier.pth', map_location=device)
classifier.load_state_dict(ckpt['model'])
classifier.eval()

provider = 'CUDAExecutionProvider' if torch.cuda.is_available() else 'CPUExecutionProvider'
app = FaceAnalysis(name='buffalo_l', providers=[provider])
app.prepare(ctx_id=0, det_size=(864,576), det_thresh=0.5)

# 얼굴 임베딩 함수
def get_emb(img):
    faces = app.get(img)
    if not faces: return None
    e = faces[0].normed_embedding
    return e / np.linalg.norm(e)

# 단일 이미지 모드
def single(img_path):
    img = cv2.imread(img_path)
    if img is None:
        print("Image load failed")
        return
    if img.shape[:2] != (576,864): img = cv2.resize(img, (864,576))
    emb = get_emb(img)
    if emb is None:
        print("No face detected")
        return
    inp = torch.tensor(emb, dtype=torch.float32).unsqueeze(0).to(device)
    out = classifier(inp)
    probs = torch.softmax(out,1)
    conf, pred = probs.max(1)
    conf = conf.item(); pid = pred.item()
    name = id2label[pid] if conf>args.threshold else 'Unknown'
    print(f"{img_path} => {name} ({conf:.2f})")

# 웹캠 모드
def webcam():
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 864)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 576)
    while True:
        ret, frame = cap.read()
        if not ret: break
        emb = get_emb(frame)
        text = 'No Face'
        if emb is not None:
            inp = torch.tensor(emb, dtype=torch.float32).unsqueeze(0).to(device)
            out = classifier(inp)
            probs = torch.softmax(out,1)
            conf, pred = probs.max(1)
            conf = conf.item(); pid = pred.item()
            text = id2label[pid] if conf>args.threshold else 'Unknown'
            text += f" ({conf:.2f})"
        cv2.putText(frame, text, (50,50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0) if 'Unknown' not in text else (0,0,255), 2)
        cv2.imshow('Recognition', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'): break
    cap.release(); cv2.destroyAllWindows()

if args.image:
    single(args.image)
else:
    webcam()