| import numpy as np |
| import cv2 |
| import torch |
| from PIL import Image |
| import os |
|
|
| from .face_recognition import Siamese, FaceClassifier, trnscm, classes |
|
|
| current_path = os.path.dirname(os.path.abspath(__file__)) |
|
|
| _SIAMESE = None |
| _CLASSIFIER = None |
| _DEVICE = None |
|
|
| def _to_bgr_np(img): |
| if img is None: |
| return None |
| if isinstance(img, str): |
| return cv2.imread(img) |
| if isinstance(img, Image.Image): |
| rgb = np.array(img.convert("RGB")) |
| return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) |
| if isinstance(img, np.ndarray): |
| arr = img |
| if arr.ndim == 2: |
| return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR) |
| if arr.ndim == 3 and arr.shape[2] == 4: |
| arr = arr[:, :, :3] |
| if arr.ndim == 3 and arr.shape[2] == 3: |
| return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) |
| return arr |
| return None |
|
|
| def _load_models(): |
| global _SIAMESE, _CLASSIFIER, _DEVICE |
| if _SIAMESE is not None and _CLASSIFIER is not None: |
| return _SIAMESE, _CLASSIFIER, _DEVICE |
|
|
| _DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
|
| siamese_path = os.path.join(current_path, "siamese_model.t7") |
| clf_path = os.path.join(current_path, "face_classifier_model.pth") |
|
|
| if not os.path.exists(siamese_path): |
| raise FileNotFoundError(f"Missing siamese model: {siamese_path}") |
| if not os.path.exists(clf_path): |
| raise FileNotFoundError(f"Missing classifier model: {clf_path}") |
|
|
| siamese = Siamese().to(_DEVICE) |
| ckpt = torch.load(siamese_path, map_location=_DEVICE) |
| if isinstance(ckpt, dict) and "net_dict" in ckpt: |
| siamese.load_state_dict(ckpt["net_dict"], strict=True) |
| else: |
| siamese.load_state_dict(ckpt, strict=True) |
| siamese.eval() |
|
|
| classifier = FaceClassifier(input_dim=5, num_classes=len(classes)).to(_DEVICE) |
| classifier.load_state_dict(torch.load(clf_path, map_location=_DEVICE), strict=True) |
| classifier.eval() |
|
|
| _SIAMESE, _CLASSIFIER = siamese, classifier |
| return _SIAMESE, _CLASSIFIER, _DEVICE |
|
|
| def detected_face(image_bgr): |
| face_haar = os.path.join(current_path, "haarcascade_frontalface_default.xml") |
| face_cascade = cv2.CascadeClassifier(face_haar) |
| if face_cascade.empty(): |
| return 0 |
|
|
| gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) |
| faces = face_cascade.detectMultiScale(gray, 1.3, 5) |
| if faces is None or len(faces) == 0: |
| return 0 |
|
|
| x, y, w, h = max(faces, key=lambda f: f[2]*f[3]) |
| crop = gray[y:y+h, x:x+w] |
| return Image.fromarray(crop) |
|
|
| def get_similarity(img1, img2): |
| siamese, classifier, device = _load_models() |
|
|
| i1 = _to_bgr_np(img1) |
| i2 = _to_bgr_np(img2) |
| if i1 is None or i2 is None: |
| return 0.0 |
|
|
| f1 = detected_face(i1) |
| f2 = detected_face(i2) |
|
|
| if f1 == 0: |
| f1 = Image.fromarray(cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY)) |
| if f2 == 0: |
| f2 = Image.fromarray(cv2.cvtColor(i2, cv2.COLOR_BGR2GRAY)) |
|
|
| t1 = trnscm(f1).unsqueeze(0).to(device) |
| t2 = trnscm(f2).unsqueeze(0).to(device) |
|
|
| with torch.no_grad(): |
| e1 = siamese.forward_once(t1) |
| e2 = siamese.forward_once(t2) |
| dist = torch.nn.functional.pairwise_distance(e1, e2).item() |
|
|
| return float(dist) |
|
|
| def get_face_class(img1): |
| siamese, classifier, device = _load_models() |
|
|
| i1 = _to_bgr_np(img1) |
| if i1 is None: |
| return "UNKNOWN" |
|
|
| f1 = detected_face(i1) |
| if f1 == 0: |
| f1 = Image.fromarray(cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY)) |
|
|
| t1 = trnscm(f1).unsqueeze(0).to(device) |
|
|
| with torch.no_grad(): |
| emb = siamese.forward_once(t1) |
| logits = classifier(emb) |
| pred = int(torch.argmax(logits, dim=1).item()) |
|
|
| if pred < 0 or pred >= len(classes): |
| return "UNKNOWN" |
| return str(classes[pred]) |
