face_utils.py

import base64
import pickle
import time
from pathlib import Path

import cv2
import numpy as np
from deepface import DeepFace
from deepface.modules.verification import find_distance, find_threshold

MODEL = "ArcFace"
METRIC = "cosine"
THRESH = find_threshold(MODEL, METRIC) * 1.12
ENROLL_DETS = ["retinaface", "mtcnn", "ssd", "opencv"]
SCAN_DETS = ["mtcnn", "retinaface", "ssd", "opencv"]


def _biggest(results):
    return max(results, key=lambda r: r["facial_area"]["w"] * r["facial_area"]["h"])


def _resize(img, max_side):
    h, w = img.shape[:2]
    if max(h, w) > max_side:
        s = max_side / max(h, w)
        return cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
    if min(h, w) < 240:
        s = 240 / min(h, w)
        return cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_CUBIC)
    return img


def _for_scan(img):
    h, w = img.shape[:2]
    if min(h, w) < 640:
        s = 640 / min(h, w)
        return cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_CUBIC)
    return img


def _to_orig(area, small, full):
    sh, sw = small.shape[:2]
    fh, fw = full.shape[:2]
    sx, sy = fw / sw, fh / sh
    return {
        "x": int(area["x"] * sx),
        "y": int(area["y"] * sy),
        "w": max(1, int(area["w"] * sx)),
        "h": max(1, int(area["h"] * sy)),
    }


def _embed(img, dets, loose=False):
    order = (False, True) if loose else (True, False)
    for enforce in order:
        for det in dets:
            try:
                out = DeepFace.represent(
                    img_path=img,
                    model_name=MODEL,
                    detector_backend=det,
                    enforce_detection=enforce,
                    align=True,
                )
                if out:
                    return out
            except (ValueError, Exception):
                pass
    return None


def encode_face_from_image(path):
    img = cv2.imread(path)
    if img is None:
        return None
    out = _embed(_resize(img, 400), ENROLL_DETS)
    return np.array(_biggest(out)["embedding"]) if out else None


def encoding_to_bytes(enc):
    return pickle.dumps((MODEL, enc))


def bytes_to_encoding(raw):
    obj = pickle.loads(raw)
    if isinstance(obj, np.ndarray):
        return None
    if isinstance(obj, dict):
        obj = (obj.get("m"), obj.get("e"))
    return np.array(obj[1]) if obj[0] == MODEL else None


def find_matching_employee(frame, known):
    if not known:
        return None

    out, scan = None, None
    for src in (
        frame,
        cv2.resize(frame, None, fx=1.6, fy=1.6, interpolation=cv2.INTER_CUBIC),
    ):
        scan = _for_scan(src)
        out = _embed(scan, SCAN_DETS) or _embed(
            scan, ["retinaface", "mtcnn"], loose=True
        )
        if out:
            break
    if not out:
        return None

    hit = _biggest(out)
    probe = np.array(hit["embedding"])
    ranked = sorted(
        (find_distance(probe, emp["enc"], METRIC), i) for i, emp in enumerate(known)
    )
    if not ranked or ranked[0][0] > THRESH:
        return None
    if len(ranked) > 1 and ranked[1][0] - ranked[0][0] < 0.06:
        return None

    emp = known[ranked[0][1]]
    return {
        "id": emp["id"],
        "name": emp["name"],
        "english_name": emp["english_name"],
        "facial_area": _to_orig(hit["facial_area"], scan, frame),
        "distance": ranked[0][0],
    }


def draw_face_box(frame, area, label=None):
    fh, fw = frame.shape[:2]
    x, y, w, h = area["x"], area["y"], area["w"], area["h"]

    pad = int(max(w, h) * 0.2)
    x = max(0, x - pad)
    y = max(0, y - pad)
    w = min(fw - x, w + pad * 2)
    h = min(fh - y, h + pad * 2)

    cv2.rectangle(frame, (x, y), (x + w, y + h), (80, 220, 120), 2)
    if not label:
        return

    txt = "".join(c for c in label if ord(c) < 128)[:32].strip() or "FACE"
    bar = max(22, min(30, h // 6))
    fs, thick = 0.65, 2
    tw, th = cv2.getTextSize(txt, cv2.FONT_HERSHEY_DUPLEX, fs, thick)[0]
    while tw > w - 12 and fs > 0.35:
        fs -= 0.05
        thick = max(1, int(fs * 2))
        tw, th = cv2.getTextSize(txt, cv2.FONT_HERSHEY_DUPLEX, fs, thick)[0]

    cv2.rectangle(frame, (x, y), (x + w, y + bar), (80, 220, 120), cv2.FILLED)
    cv2.putText(
        frame,
        txt,
        (x + max(4, (w - tw) // 2), y + bar - max(5, (bar - th) // 2)),
        cv2.FONT_HERSHEY_DUPLEX,
        fs,
        (20, 30, 20),
        thick,
        cv2.LINE_AA,
    )


def save_base64_photo(data_url, dest):
    if not data_url or "," not in data_url:
        return None

    hdr, b64 = data_url.split(",", 1)
    if "image" not in hdr:
        return None

    try:
        raw = base64.b64decode(b64)
    except (ValueError, TypeError):
        return None
    if len(raw) < 1000:
        return None

    dest.mkdir(parents=True, exist_ok=True)
    path = dest / f"capture_{int(time.time())}.jpg"
    n = 1
    while path.exists():
        path = dest / f"capture_{int(time.time())}_{n}.jpg"
        n += 1

    path.write_bytes(raw)
    if cv2.imread(str(path)) is None:
        path.unlink(missing_ok=True)
        return None
    return path.name


def save_upload(f, dest):
    dest.mkdir(parents=True, exist_ok=True)
    name = f.filename or "photo.jpg"
    safe = "".join(c if c.isalnum() or c in "._-" else "_" for c in name)
    path = dest / safe

    n = 1
    while path.exists():
        stem, suf = Path(safe).stem, Path(safe).suffix or ".jpg"
        path = dest / f"{stem}_{n}{suf}"
        n += 1

    f.save(path)
    return path.name