File size: 3,796 Bytes
820b49d 9040d50 62c391f 820b49d a150f43 820b49d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | 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]) |