126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from werkzeug.datastructures import FileStorage
|
|
|
|
from app.models import Person
|
|
|
|
|
|
class RecognitionUnavailable(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DetectedFace:
|
|
person: "Person | None"
|
|
confidence: float | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RecognitionResult:
|
|
faces: list[DetectedFace]
|
|
annotated_image: str
|
|
|
|
|
|
def encode_uploaded_image(file: "FileStorage") -> list[float]:
|
|
cv2, np, face_recognition = _load_dependencies()
|
|
raw = file.read()
|
|
image = _decode_bytes(raw, cv2, np)
|
|
rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
encodings = face_recognition.face_encodings(rgb)
|
|
if len(encodings) == 0:
|
|
raise ValueError(f"No face was detected in {file.filename}.")
|
|
if len(encodings) > 1:
|
|
raise ValueError(f"Upload a single-face photo for {file.filename}.")
|
|
return [float(value) for value in encodings[0]]
|
|
|
|
|
|
def recognize_faces(data_url: str, people: list["Person"], *, tolerance: float) -> RecognitionResult:
|
|
cv2, np, face_recognition = _load_dependencies()
|
|
image = _decode_data_url(data_url, cv2, np)
|
|
rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
locations = face_recognition.face_locations(rgb)
|
|
encodings = face_recognition.face_encodings(rgb, locations)
|
|
candidates = [
|
|
(person, encoding.vector)
|
|
for person in people
|
|
for encoding in person.face_encodings
|
|
]
|
|
|
|
faces: list[DetectedFace] = []
|
|
for location, face_encoding in zip(locations, encodings, strict=False):
|
|
person = None
|
|
confidence = None
|
|
label = "Unknown"
|
|
|
|
if candidates:
|
|
vectors = [candidate[1] for candidate in candidates]
|
|
distances = face_recognition.face_distance(vectors, face_encoding)
|
|
best_index = int(np.argmin(distances))
|
|
best_distance = float(distances[best_index])
|
|
if best_distance <= tolerance:
|
|
person = candidates[best_index][0]
|
|
confidence = max(0.0, 1.0 - best_distance)
|
|
label = person.full_name
|
|
|
|
faces.append(DetectedFace(person=person, confidence=confidence))
|
|
_draw_face_label(image, location, label, cv2)
|
|
|
|
if not faces:
|
|
return RecognitionResult(faces=[DetectedFace(person=None, confidence=None)], annotated_image=_encode_png(image, cv2))
|
|
return RecognitionResult(faces=faces, annotated_image=_encode_png(image, cv2))
|
|
|
|
|
|
def _load_dependencies():
|
|
try:
|
|
import cv2
|
|
import face_recognition
|
|
import numpy as np
|
|
except ImportError as exc:
|
|
raise RecognitionUnavailable(
|
|
"Face recognition dependencies are not installed. Install requirements on a supported Python version."
|
|
) from exc
|
|
return cv2, np, face_recognition
|
|
|
|
|
|
def _decode_data_url(data_url: str, cv2, np):
|
|
if "," not in data_url:
|
|
raise ValueError("Invalid image payload.")
|
|
header, encoded = data_url.split(",", 1)
|
|
if not header.startswith("data:image/"):
|
|
raise ValueError("Only image uploads are supported.")
|
|
try:
|
|
raw = base64.b64decode(encoded, validate=True)
|
|
except ValueError as exc:
|
|
raise ValueError("Invalid image encoding.") from exc
|
|
return _decode_bytes(raw, cv2, np)
|
|
|
|
|
|
def _decode_bytes(raw: bytes, cv2, np):
|
|
if not raw:
|
|
raise ValueError("The uploaded image is empty.")
|
|
array = np.frombuffer(raw, np.uint8)
|
|
image = cv2.imdecode(array, cv2.IMREAD_COLOR)
|
|
if image is None:
|
|
raise ValueError("The uploaded file is not a readable image.")
|
|
return image
|
|
|
|
|
|
def _draw_face_label(image, location, label: str, cv2) -> None:
|
|
top, right, bottom, left = location
|
|
cv2.rectangle(image, (left, top), (right, bottom), (30, 136, 229), 2)
|
|
cv2.rectangle(image, (left, bottom - 28), (right, bottom), (30, 136, 229), cv2.FILLED)
|
|
cv2.putText(image, label[:32], (left + 6, bottom - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1)
|
|
|
|
|
|
def _encode_png(image, cv2) -> str:
|
|
ok, buffer = cv2.imencode(".png", image)
|
|
if not ok:
|
|
raise ValueError("Could not encode processed image.")
|
|
return "data:image/png;base64," + base64.b64encode(buffer).decode("ascii")
|