Spaces:
Runtime error
Runtime error
| """ | |
| Face pipeline for the live demo: detect -> crop -> preprocess -> embed -> match. | |
| The one rule this module exists to enforce: enrollment and recognition run the | |
| EXACT same code path. ArcFace embeddings are only comparable when the crop | |
| geometry and lighting normalisation match, so both directions call | |
| extract_faces() + embed() and nothing else. | |
| Galleries are per-session and in-memory. Nothing is written to disk on the | |
| server: the user's .pkl belongs to the user's device. The production | |
| student_embeddings.pkl is never read or written here. | |
| """ | |
| import base64 | |
| import io | |
| import pickle | |
| import threading | |
| import cv2 | |
| import numpy as np | |
| from numpy.linalg import norm | |
| MODEL_NAME = "ArcFace" | |
| PAD_RATIO = 0.20 # matches main.py: 20% padding around the MTCNN box | |
| DETECT_MAX_DIM = 640 # downscale before detection; boxes scaled back up | |
| MIN_FACE_CONFIDENCE = 0.90 | |
| # TensorFlow graphs are not safe to call from multiple threads at once and | |
| # Flask's dev server is threaded, so every inference goes through this lock. | |
| _LOCK = threading.Lock() | |
| _detector = None | |
| # -------------------------------------------------------------------------- | |
| # model loading | |
| # -------------------------------------------------------------------------- | |
| def get_detector(): | |
| global _detector | |
| if _detector is None: | |
| from mtcnn import MTCNN | |
| _detector = MTCNN() | |
| return _detector | |
| def warmup(): | |
| """Load MTCNN + ArcFace and run one dummy inference of each. | |
| Called at server start so the first real request isn't a 30s stall. | |
| """ | |
| from deepface import DeepFace | |
| dummy = np.zeros((160, 160, 3), dtype=np.uint8) | |
| with _LOCK: | |
| get_detector().detect_faces(cv2.cvtColor(dummy, cv2.COLOR_BGR2RGB)) | |
| DeepFace.represent( | |
| img_path=dummy, | |
| model_name=MODEL_NAME, | |
| detector_backend="skip", | |
| enforce_detection=False, | |
| ) | |
| # -------------------------------------------------------------------------- | |
| # image codec helpers | |
| # -------------------------------------------------------------------------- | |
| def decode_data_url(data_url): | |
| """A data: URL or bare base64 string -> BGR ndarray.""" | |
| if "," in data_url: | |
| data_url = data_url.split(",", 1)[1] | |
| raw = base64.b64decode(data_url) | |
| buf = np.frombuffer(raw, dtype=np.uint8) | |
| img = cv2.imdecode(buf, cv2.IMREAD_COLOR) | |
| if img is None: | |
| raise ValueError("could not decode image") | |
| return img | |
| def encode_jpeg_data_url(bgr, quality=85): | |
| ok, buf = cv2.imencode(".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), quality]) | |
| if not ok: | |
| raise ValueError("could not encode image") | |
| return "data:image/jpeg;base64," + base64.b64encode(buf.tobytes()).decode("ascii") | |
| def l2_normalize(x): | |
| x = np.asarray(x, dtype=np.float32) | |
| return x / (norm(x) + 1e-10) | |
| def cosine_similarity(a, b): | |
| return float(np.dot(a, b) / (norm(a) * norm(b) + 1e-10)) | |
| # -------------------------------------------------------------------------- | |
| # detection + preprocessing (same logic as main.py's preprocess_face) | |
| # -------------------------------------------------------------------------- | |
| def preprocess_face(face_rgb, keypoints=None, origin=None): | |
| """Adaptive gamma -> eye-alignment rotation -> Y-channel histogram EQ.""" | |
| try: | |
| if face_rgb is None or face_rgb.size == 0: | |
| return None | |
| gray = cv2.cvtColor(face_rgb, cv2.COLOR_RGB2GRAY) | |
| mean = float(np.mean(gray)) | |
| gamma = 1.0 | |
| if mean < 100: | |
| gamma = 1.5 | |
| elif mean > 160: | |
| gamma = 0.7 | |
| lut = np.array( | |
| [((i / 255.0) ** (1.0 / gamma)) * 255 for i in np.arange(256)] | |
| ).astype("uint8") | |
| face_rgb = cv2.LUT(face_rgb, lut) | |
| if keypoints and "left_eye" in keypoints and "right_eye" in keypoints: | |
| le, re = keypoints["left_eye"], keypoints["right_eye"] | |
| if le is not None and re is not None and len(le) == 2 and len(re) == 2: | |
| lx, ly = float(le[0]), float(le[1]) | |
| rx, ry = float(re[0]), float(re[1]) | |
| if not any(np.isnan(v) for v in (lx, ly, rx, ry)): | |
| if origin is not None: | |
| ox, oy = origin | |
| lx -= ox | |
| rx -= ox | |
| ly -= oy | |
| ry -= oy | |
| angle = float(np.degrees(np.arctan2(ry - ly, rx - lx))) | |
| h, w = face_rgb.shape[:2] | |
| cx = max(0.0, min(float(w - 1), (lx + rx) / 2.0)) | |
| cy = max(0.0, min(float(h - 1), (ly + ry) / 2.0)) | |
| m = cv2.getRotationMatrix2D((cx, cy), angle, 1.0) | |
| face_rgb = cv2.warpAffine( | |
| face_rgb, m, (w, h), | |
| flags=cv2.INTER_LINEAR, | |
| borderMode=cv2.BORDER_REPLICATE, | |
| ) | |
| yuv = cv2.cvtColor(face_rgb, cv2.COLOR_RGB2YUV) | |
| yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0]) | |
| return cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB) | |
| except Exception: | |
| return face_rgb | |
| def extract_faces(bgr): | |
| """Detect every face in a BGR frame. | |
| Detection runs on a downscaled copy (MTCNN cost scales with pixel count) | |
| but crops are taken from the full-resolution original so no detail is lost. | |
| Returns list of {box, confidence, crop_bgr} sorted left-to-right. | |
| """ | |
| rgb_full = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) | |
| h, w = rgb_full.shape[:2] | |
| scale = 1.0 | |
| rgb_small = rgb_full | |
| if max(h, w) > DETECT_MAX_DIM: | |
| scale = DETECT_MAX_DIM / float(max(h, w)) | |
| rgb_small = cv2.resize( | |
| rgb_full, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA | |
| ) | |
| with _LOCK: | |
| results = get_detector().detect_faces(rgb_small) | |
| faces = [] | |
| for res in results: | |
| if res.get("confidence", 0.0) < MIN_FACE_CONFIDENCE: | |
| continue | |
| try: | |
| x, y, bw, bh = [int(round(v / scale)) for v in res["box"]] | |
| except Exception: | |
| continue | |
| if bw <= 0 or bh <= 0: | |
| continue | |
| pad = int(max(1, min(bw, bh) * PAD_RATIO)) | |
| x1, y1 = max(0, x - pad), max(0, y - pad) | |
| x2, y2 = min(w, x + bw + pad), min(h, y + bh + pad) | |
| if x2 <= x1 or y2 <= y1: | |
| continue | |
| crop = rgb_full[y1:y2, x1:x2] | |
| if crop.size == 0: | |
| continue | |
| kps = res.get("keypoints") or {} | |
| kps = {k: (v[0] / scale, v[1] / scale) for k, v in kps.items() if v is not None} | |
| processed = preprocess_face(crop, kps, origin=(x1, y1)) | |
| if processed is None or processed.size == 0: | |
| continue | |
| faces.append({ | |
| "box": [int(x), int(y), int(bw), int(bh)], | |
| "confidence": float(res.get("confidence", 0.0)), | |
| "crop_bgr": cv2.cvtColor(processed, cv2.COLOR_RGB2BGR), | |
| }) | |
| faces.sort(key=lambda f: f["box"][0]) | |
| return faces | |
| def embed(crop_bgr): | |
| """512-d L2-normalised ArcFace vector. Array in, no disk round-trip.""" | |
| from deepface import DeepFace | |
| with _LOCK: | |
| rep = DeepFace.represent( | |
| img_path=crop_bgr, | |
| model_name=MODEL_NAME, | |
| detector_backend="skip", | |
| enforce_detection=False, | |
| ) | |
| return l2_normalize(rep[0]["embedding"]) | |
| # -------------------------------------------------------------------------- | |
| # gallery | |
| # -------------------------------------------------------------------------- | |
| def cohesion(embs): | |
| """Mean pairwise similarity within one person's own embeddings. | |
| Enrollment quality signal: low cohesion means the reference photos | |
| disagree with each other, and matching against them will be unreliable. | |
| """ | |
| if len(embs) < 2: | |
| return None | |
| a = np.asarray(embs, dtype=np.float32) | |
| sims = a @ a.T | |
| n = len(a) | |
| return float((sims.sum() - np.trace(sims)) / (n * (n - 1))) | |
| class RestrictedUnpickler(pickle.Unpickler): | |
| """pickle.load() on user-supplied bytes is arbitrary code execution. | |
| A .pkl uploaded from a browser is untrusted input, so imports are limited | |
| to what a numpy embedding file legitimately needs. Anything else raises | |
| instead of executing. | |
| """ | |
| _ALLOWED = { | |
| ("numpy", "ndarray"), | |
| ("numpy", "dtype"), | |
| ("numpy", "float32"), | |
| ("numpy", "float64"), | |
| ("numpy.core.multiarray", "_reconstruct"), | |
| ("numpy.core.multiarray", "scalar"), | |
| ("numpy._core.multiarray", "_reconstruct"), | |
| ("numpy._core.multiarray", "scalar"), | |
| } | |
| def find_class(self, module, name): | |
| if (module, name) in self._ALLOWED: | |
| return super().find_class(module, name) | |
| raise pickle.UnpicklingError(f"blocked for safety: {module}.{name}") | |
| def restricted_loads(raw): | |
| return RestrictedUnpickler(io.BytesIO(raw)).load() | |
| class Gallery: | |
| """{person_id: [embedding, ...]} with the same match rule as main.py: | |
| a person's score is the MAX similarity over their enrolled embeddings. | |
| Purely in-memory and scoped to one browser session. Persistence is the | |
| client's job -- see to_pickle_bytes() / to_json(). | |
| """ | |
| def __init__(self, data=None): | |
| self.data = data or {} | |
| # ---- serialisation ------------------------------------------------- | |
| def to_pickle_bytes(self): | |
| """Byte-identical in structure to what face_embd.py writes, so an | |
| exported file drops straight into main.py with no conversion.""" | |
| return pickle.dumps( | |
| {k: [np.asarray(e, dtype=np.float32) for e in v] | |
| for k, v in self.data.items()}, | |
| protocol=4, | |
| ) | |
| def to_json(self): | |
| return {k: [[float(x) for x in e] for e in v] for k, v in self.data.items()} | |
| def _validate(obj): | |
| if not isinstance(obj, dict): | |
| raise ValueError("expected a dict of {person_id: [embedding, ...]}") | |
| clean = {} | |
| for pid, embs in obj.items(): | |
| if not isinstance(pid, str) or not pid.strip(): | |
| raise ValueError("person ids must be non-empty strings") | |
| if not isinstance(embs, (list, tuple)) or not embs: | |
| raise ValueError(f"'{pid}' has no embeddings") | |
| out = [] | |
| for e in embs: | |
| a = np.asarray(e, dtype=np.float32).ravel() | |
| if a.size != 512 or not np.all(np.isfinite(a)): | |
| raise ValueError(f"'{pid}' has a bad embedding (size {a.size})") | |
| out.append(l2_normalize(a)) | |
| clean[pid.strip()] = out | |
| return clean | |
| def from_json(cls, obj): | |
| return cls(cls._validate(obj)) | |
| def from_pickle_bytes(cls, raw): | |
| return cls(cls._validate(restricted_loads(raw))) | |
| # ---- gallery ops --------------------------------------------------- | |
| def add(self, person_id, embedding): | |
| self.data.setdefault(person_id, []).append(embedding) | |
| def remove(self, person_id): | |
| return self.data.pop(person_id, None) is not None | |
| def clear(self): | |
| self.data = {} | |
| def people(self): | |
| return sorted(self.data.keys()) | |
| def summary(self): | |
| return [ | |
| {"id": pid, "count": len(embs), "cohesion": cohesion(embs)} | |
| for pid, embs in sorted(self.data.items()) | |
| ] | |
| def merge(self, source, target): | |
| """Fold source's embeddings into target and drop the source name. | |
| Used when two names turn out to be one face. Stored embeddings are | |
| moved as-is, never recomputed, so the merged identity is exactly the | |
| union of both registrations. | |
| """ | |
| if not source or not target or source == target: | |
| raise ValueError("pick two different people to merge") | |
| if source not in self.data or target not in self.data: | |
| raise ValueError("both people must be registered in this session") | |
| self.data[target].extend(self.data.pop(source)) | |
| def rank(self, probe): | |
| """All enrolled people scored against the probe, best first.""" | |
| scored = [ | |
| {"id": pid, "score": max(cosine_similarity(probe, e) for e in embs)} | |
| for pid, embs in self.data.items() | |
| if embs | |
| ] | |
| scored.sort(key=lambda s: s["score"], reverse=True) | |
| return scored | |
| # -------------------------------------------------------------------------- | |
| # duplicate identities | |
| # | |
| # Pure analysis over embeddings the pipeline has already produced. Nothing | |
| # here touches detection, preprocessing, embedding or rank(). | |
| # -------------------------------------------------------------------------- | |
| MATCH_THRESHOLD = 0.45 # main.py's recognition threshold | |
| # Calibrated on raw_dataset (9 students). Each student's photos split in half | |
| # stood in for one person under two names (540 splits); every pair of different | |
| # students stood in for different people (576 comparisons). This combined rule | |
| # caught 98.5% of the same-person splits and flagged 0 of the different pairs. | |
| DUP_RATIO_THRESHOLD = 0.88 | |
| DUP_CENTROID_FLOOR = 0.60 | |
| COHESION_FLOOR = 0.30 # stops a scattered registration from inflating the ratio | |
| def compare_identities(a, b): | |
| """How alike two registered names are, from their stored embeddings. | |
| ratio - mean cross-name similarity divided by each name's own | |
| cohesion: "are B's photos as close to A's as A's photos are to | |
| each other?" ~1.0 means indistinguishable. The normalisation is | |
| what makes this work: in this pipeline even DIFFERENT students | |
| have raw cross-similarity around 0.3-0.5, so no absolute cosine | |
| cutoff separates them. | |
| centroid - cosine between the two names' mean embeddings. An independent | |
| guard for the ratio's weak spot (dividing by a low cohesion). | |
| match_rate - main.py's rule turned on the gallery: the share of each name's | |
| photos recognised as the other at 0.45. Reported only -- it is | |
| not diagnostic, different students often score high on it too. | |
| """ | |
| A = np.asarray(a, dtype=np.float32) | |
| B = np.asarray(b, dtype=np.float32) | |
| S = A @ B.T | |
| cross = float(S.mean()) | |
| ca, cb = A.mean(axis=0), B.mean(axis=0) | |
| centroid = float(ca @ cb / (norm(ca) * norm(cb) + 1e-10)) | |
| ka, kb = cohesion(A), cohesion(B) | |
| ka = max(ka if ka is not None else cross, COHESION_FLOOR) | |
| kb = max(kb if kb is not None else cross, COHESION_FLOOR) | |
| hits = int((S.max(axis=1) >= MATCH_THRESHOLD).sum() | |
| + (S.max(axis=0) >= MATCH_THRESHOLD).sum()) | |
| return { | |
| "ratio": cross / float(np.sqrt(ka * kb)), | |
| "centroid": centroid, | |
| "cross_mean": cross, | |
| "match_rate": hits / float(len(A) + len(B)), | |
| } | |
| def find_duplicates(data, dismissed=()): | |
| """Pairs of different names whose faces look like the same person. | |
| dismissed holds frozenset({a, b}) pairs the user has confirmed are | |
| different people, so they aren't raised again. | |
| """ | |
| ids = sorted(k for k, v in data.items() if v) | |
| out = [] | |
| for i, a in enumerate(ids): | |
| for b in ids[i + 1:]: | |
| if frozenset((a, b)) in dismissed: | |
| continue | |
| m = compare_identities(data[a], data[b]) | |
| if m["ratio"] >= DUP_RATIO_THRESHOLD and m["centroid"] >= DUP_CENTROID_FLOOR: | |
| out.append({"a": a, "b": b, **{k: round(v, 4) for k, v in m.items()}}) | |
| out.sort(key=lambda d: d["ratio"], reverse=True) | |
| return out | |