mirror of
https://github.com/yakhyo/uniface.git
synced 2026-05-14 20:35:28 +00:00
Compare commits
7 Commits
feat/unifa
...
v3.3.0
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
426bd71505 | ||
|
|
ede8b27091 | ||
|
|
02c77ce5db | ||
|
|
d70d6a254f | ||
|
|
7d37633b1a | ||
|
|
bc413df4a8 | ||
|
|
8db0577991 |
6
.github/workflows/ci.yml
vendored
6
.github/workflows/ci.yml
vendored
@@ -20,7 +20,7 @@ jobs:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.10"
|
||||
python-version: "3.11"
|
||||
- uses: pre-commit/action@v3.0.1
|
||||
|
||||
test:
|
||||
@@ -35,8 +35,12 @@ jobs:
|
||||
# Full Python range on Linux (fastest runner)
|
||||
- os: ubuntu-latest
|
||||
python-version: "3.10"
|
||||
- os: ubuntu-latest
|
||||
python-version: "3.11"
|
||||
- os: ubuntu-latest
|
||||
python-version: "3.13"
|
||||
- os: ubuntu-latest
|
||||
python-version: "3.14"
|
||||
- os: macos-latest
|
||||
python-version: "3.13"
|
||||
- os: windows-latest
|
||||
|
||||
3
.github/workflows/docs.yml
vendored
3
.github/workflows/docs.yml
vendored
@@ -2,7 +2,8 @@ name: Deploy docs
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
tags:
|
||||
- "v*.*.*"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
|
||||
4
.github/workflows/publish.yml
vendored
4
.github/workflows/publish.yml
vendored
@@ -54,7 +54,7 @@ jobs:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python-version: ["3.10", "3.13"]
|
||||
python-version: ["3.10", "3.11", "3.13"]
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
@@ -92,7 +92,7 @@ jobs:
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.10"
|
||||
python-version: "3.11"
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install build tools
|
||||
|
||||
@@ -59,12 +59,12 @@ This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and formattin
|
||||
#### General Rules
|
||||
|
||||
- **Line length:** 120 characters maximum
|
||||
- **Python version:** 3.10+ (use modern syntax)
|
||||
- **Python version:** 3.11+ (use modern syntax)
|
||||
- **Quote style:** Single quotes for strings, double quotes for docstrings
|
||||
|
||||
#### Type Hints
|
||||
|
||||
Use modern Python 3.10+ type hints (PEP 585 and PEP 604):
|
||||
Use modern Python 3.11+ type hints (PEP 585 and PEP 604):
|
||||
|
||||
```python
|
||||
# Preferred (modern)
|
||||
@@ -184,6 +184,9 @@ Example notebooks demonstrating library usage:
|
||||
| Face Parsing | [06_face_parsing.ipynb](examples/06_face_parsing.ipynb) |
|
||||
| Face Anonymization | [07_face_anonymization.ipynb](examples/07_face_anonymization.ipynb) |
|
||||
| Gaze Estimation | [08_gaze_estimation.ipynb](examples/08_gaze_estimation.ipynb) |
|
||||
| Face Segmentation | [09_face_segmentation.ipynb](examples/09_face_segmentation.ipynb) |
|
||||
| Face Vector Store | [10_face_vector_store.ipynb](examples/10_face_vector_store.ipynb) |
|
||||
| Head Pose Estimation | [11_head_pose_estimation.ipynb](examples/11_head_pose_estimation.ipynb) |
|
||||
|
||||
## Questions?
|
||||
|
||||
|
||||
81
README.md
81
README.md
@@ -31,8 +31,9 @@
|
||||
- **Facial Landmarks** — 106-point landmark localization module (separate from 5-point detector landmarks)
|
||||
- **Face Parsing** — BiSeNet semantic segmentation (19 classes), XSeg face masking
|
||||
- **Gaze Estimation** — Real-time gaze direction with MobileGaze
|
||||
- **Head Pose Estimation** — 3D head orientation (pitch, yaw, roll) with 6D rotation representation
|
||||
- **Attribute Analysis** — Age, gender, race (FairFace), and emotion
|
||||
- **Vector Indexing** — FAISS-backed embedding store for fast multi-identity search
|
||||
- **Vector Store** — FAISS-backed embedding store for fast multi-identity search
|
||||
- **Anti-Spoofing** — Face liveness detection with MiniFASNet
|
||||
- **Face Anonymization** — 5 blur methods for privacy protection
|
||||
- **Hardware Acceleration** — ARM64 (Apple Silicon), CUDA (NVIDIA), CPU
|
||||
@@ -60,7 +61,7 @@ git clone https://github.com/yakhyo/uniface.git
|
||||
cd uniface && pip install -e .
|
||||
```
|
||||
|
||||
**FAISS vector indexing**
|
||||
**FAISS vector store**
|
||||
|
||||
```bash
|
||||
pip install faiss-cpu # or faiss-gpu for CUDA
|
||||
@@ -126,14 +127,10 @@ for face in faces:
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from uniface.analyzer import FaceAnalyzer
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.recognition import ArcFace
|
||||
from uniface import FaceAnalyzer
|
||||
|
||||
detector = RetinaFace()
|
||||
recognizer = ArcFace()
|
||||
|
||||
analyzer = FaceAnalyzer(detector, recognizer=recognizer)
|
||||
# Zero-config: uses SCRFD (500M) + ArcFace (MobileNet) by default
|
||||
analyzer = FaceAnalyzer()
|
||||
|
||||
image = cv2.imread("photo.jpg")
|
||||
if image is None:
|
||||
@@ -145,19 +142,36 @@ for face in faces:
|
||||
print(face.bbox, face.embedding.shape if face.embedding is not None else None)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Execution Providers (ONNX Runtime)
|
||||
With attributes:
|
||||
|
||||
```python
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface import FaceAnalyzer, AgeGender
|
||||
|
||||
# Force CPU-only inference
|
||||
detector = RetinaFace(providers=["CPUExecutionProvider"])
|
||||
analyzer = FaceAnalyzer(attributes=[AgeGender()])
|
||||
faces = analyzer.analyze(image)
|
||||
|
||||
for face in faces:
|
||||
print(f"{face.sex}, {face.age}y, embedding={face.embedding.shape}")
|
||||
```
|
||||
|
||||
See more in the docs:
|
||||
https://yakhyo.github.io/uniface/concepts/execution-providers/
|
||||
---
|
||||
|
||||
## Jupyter Notebooks
|
||||
|
||||
| Example | Colab | Description |
|
||||
|---------|:-----:|-------------|
|
||||
| [01_face_detection.ipynb](examples/01_face_detection.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/01_face_detection.ipynb) | Face detection and landmarks |
|
||||
| [02_face_alignment.ipynb](examples/02_face_alignment.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/02_face_alignment.ipynb) | Face alignment for recognition |
|
||||
| [03_face_verification.ipynb](examples/03_face_verification.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/03_face_verification.ipynb) | Compare faces for identity |
|
||||
| [04_face_search.ipynb](examples/04_face_search.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/04_face_search.ipynb) | Find a person in group photos |
|
||||
| [05_face_analyzer.ipynb](examples/05_face_analyzer.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/05_face_analyzer.ipynb) | All-in-one analysis |
|
||||
| [06_face_parsing.ipynb](examples/06_face_parsing.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/06_face_parsing.ipynb) | Semantic face segmentation |
|
||||
| [07_face_anonymization.ipynb](examples/07_face_anonymization.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/07_face_anonymization.ipynb) | Privacy-preserving blur |
|
||||
| [08_gaze_estimation.ipynb](examples/08_gaze_estimation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/08_gaze_estimation.ipynb) | Gaze direction estimation |
|
||||
| [09_face_segmentation.ipynb](examples/09_face_segmentation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/09_face_segmentation.ipynb) | Face segmentation with XSeg |
|
||||
| [10_face_vector_store.ipynb](examples/10_face_vector_store.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/10_face_vector_store.ipynb) | FAISS-backed face database |
|
||||
| [11_head_pose_estimation.ipynb](examples/11_head_pose_estimation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/11_head_pose_estimation.ipynb) | Head pose estimation (pitch, yaw, roll) |
|
||||
| [12_face_recognition.ipynb](examples/12_face_recognition.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/12_face_recognition.ipynb) | Standalone face recognition pipeline |
|
||||
|
||||
---
|
||||
|
||||
@@ -176,6 +190,20 @@ Full documentation: https://yakhyo.github.io/uniface/
|
||||
|
||||
---
|
||||
|
||||
## Execution Providers (ONNX Runtime)
|
||||
|
||||
```python
|
||||
from uniface.detection import RetinaFace
|
||||
|
||||
# Force CPU-only inference
|
||||
detector = RetinaFace(providers=["CPUExecutionProvider"])
|
||||
```
|
||||
|
||||
See more in the docs:
|
||||
https://yakhyo.github.io/uniface/concepts/execution-providers/
|
||||
|
||||
---
|
||||
|
||||
## Datasets
|
||||
|
||||
| Task | Training Dataset | Models |
|
||||
@@ -185,6 +213,7 @@ Full documentation: https://yakhyo.github.io/uniface/
|
||||
| Recognition | WebFace600K | ArcFace |
|
||||
| Recognition | WebFace4M / 12M | AdaFace |
|
||||
| Gaze | Gaze360 | MobileGaze |
|
||||
| Head Pose | 300W-LP | HeadPose (ResNet, MobileNet) |
|
||||
| Parsing | CelebAMask-HQ | BiSeNet |
|
||||
| Attributes | CelebA, FairFace, AffectNet | AgeGender, FairFace, Emotion |
|
||||
|
||||
@@ -192,23 +221,6 @@ Full documentation: https://yakhyo.github.io/uniface/
|
||||
|
||||
---
|
||||
|
||||
## Jupyter Notebooks
|
||||
|
||||
| Example | Colab | Description |
|
||||
|---------|:-----:|-------------|
|
||||
| [01_face_detection.ipynb](examples/01_face_detection.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/01_face_detection.ipynb) | Face detection and landmarks |
|
||||
| [02_face_alignment.ipynb](examples/02_face_alignment.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/02_face_alignment.ipynb) | Face alignment for recognition |
|
||||
| [03_face_verification.ipynb](examples/03_face_verification.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/03_face_verification.ipynb) | Compare faces for identity |
|
||||
| [04_face_search.ipynb](examples/04_face_search.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/04_face_search.ipynb) | Find a person in group photos |
|
||||
| [05_face_analyzer.ipynb](examples/05_face_analyzer.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/05_face_analyzer.ipynb) | All-in-one analysis |
|
||||
| [06_face_parsing.ipynb](examples/06_face_parsing.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/06_face_parsing.ipynb) | Semantic face segmentation |
|
||||
| [07_face_anonymization.ipynb](examples/07_face_anonymization.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/07_face_anonymization.ipynb) | Privacy-preserving blur |
|
||||
| [08_gaze_estimation.ipynb](examples/08_gaze_estimation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/08_gaze_estimation.ipynb) | Gaze direction estimation |
|
||||
| [09_face_segmentation.ipynb](examples/09_face_segmentation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/09_face_segmentation.ipynb) | Face segmentation with XSeg |
|
||||
| [10_face_vector_store.ipynb](examples/10_face_vector_store.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/10_face_vector_store.ipynb) | FAISS-backed face database |
|
||||
|
||||
---
|
||||
|
||||
## Licensing and Model Usage
|
||||
|
||||
UniFace is MIT-licensed, but several pretrained models carry their own licenses.
|
||||
@@ -234,6 +246,7 @@ If you plan commercial use, verify model license compatibility.
|
||||
| Parsing | [face-parsing](https://github.com/yakhyo/face-parsing) | ✓ | BiSeNet Face Parsing |
|
||||
| Parsing | [face-segmentation](https://github.com/yakhyo/face-segmentation) | - | XSeg Face Segmentation |
|
||||
| Gaze | [gaze-estimation](https://github.com/yakhyo/gaze-estimation) | ✓ | MobileGaze Training |
|
||||
| Head Pose | [head-pose-estimation](https://github.com/yakhyo/head-pose-estimation) | ✓ | Head Pose Training (6DRepNet-style) |
|
||||
| Anti-Spoofing | [face-anti-spoofing](https://github.com/yakhyo/face-anti-spoofing) | - | MiniFASNet Inference |
|
||||
| Attributes | [fairface-onnx](https://github.com/yakhyo/fairface-onnx) | - | FairFace ONNX Inference |
|
||||
|
||||
|
||||
BIN
assets/test_images/image5.jpg
Normal file
BIN
assets/test_images/image5.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 5.8 KiB |
@@ -39,16 +39,20 @@ recognizer = ArcFace(providers=['CPUExecutionProvider'])
|
||||
detector = RetinaFace(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
||||
```
|
||||
|
||||
All model classes accept the `providers` parameter:
|
||||
All **ONNX-based** model classes accept the `providers` parameter:
|
||||
|
||||
- Detection: `RetinaFace`, `SCRFD`, `YOLOv5Face`, `YOLOv8Face`
|
||||
- Recognition: `ArcFace`, `AdaFace`, `MobileFace`, `SphereFace`
|
||||
- Landmarks: `Landmark106`
|
||||
- Gaze: `MobileGaze`
|
||||
- Parsing: `BiSeNet`
|
||||
- Parsing: `BiSeNet`, `XSeg`
|
||||
- Attributes: `AgeGender`, `FairFace`
|
||||
- Anti-Spoofing: `MiniFASNet`
|
||||
|
||||
!!! note "Non-ONNX components"
|
||||
- **Emotion** uses TorchScript and selects its device automatically (`mps` / `cuda` / `cpu`). It does **not** accept the `providers` parameter.
|
||||
- **BlurFace** is a pure OpenCV utility and does not load any model.
|
||||
|
||||
---
|
||||
|
||||
## Check Available Providers
|
||||
|
||||
@@ -106,6 +106,27 @@ print(f"Yaw: {np.degrees(result.yaw):.1f}°")
|
||||
|
||||
---
|
||||
|
||||
### HeadPoseResult
|
||||
|
||||
```python
|
||||
@dataclass(frozen=True)
|
||||
class HeadPoseResult:
|
||||
pitch: float # Rotation around X-axis (degrees), + = looking down
|
||||
yaw: float # Rotation around Y-axis (degrees), + = looking right
|
||||
roll: float # Rotation around Z-axis (degrees), + = tilting clockwise
|
||||
```
|
||||
|
||||
**Usage:**
|
||||
|
||||
```python
|
||||
result = head_pose.estimate(face_crop)
|
||||
print(f"Pitch: {result.pitch:.1f}°")
|
||||
print(f"Yaw: {result.yaw:.1f}°")
|
||||
print(f"Roll: {result.roll:.1f}°")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### SpoofingResult
|
||||
|
||||
```python
|
||||
@@ -144,11 +165,11 @@ class AttributeResult:
|
||||
|
||||
```python
|
||||
# AgeGender model
|
||||
result = age_gender.predict(image, face.bbox)
|
||||
result = age_gender.predict(image, face)
|
||||
print(f"{result.sex}, {result.age} years old")
|
||||
|
||||
# FairFace model
|
||||
result = fairface.predict(image, face.bbox)
|
||||
result = fairface.predict(image, face)
|
||||
print(f"{result.sex}, {result.age_group}, {result.race}")
|
||||
```
|
||||
|
||||
@@ -171,7 +192,7 @@ Face recognition models return normalized 512-dimensional embeddings:
|
||||
|
||||
```python
|
||||
embedding = recognizer.get_normalized_embedding(image, landmarks)
|
||||
print(f"Shape: {embedding.shape}") # (1, 512)
|
||||
print(f"Shape: {embedding.shape}") # (512,)
|
||||
print(f"Norm: {np.linalg.norm(embedding):.4f}") # ~1.0
|
||||
```
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ graph TB
|
||||
LMK[Landmarks]
|
||||
ATTR[Attributes]
|
||||
GAZE[Gaze]
|
||||
HPOSE[Head Pose]
|
||||
PARSE[Parsing]
|
||||
SPOOF[Anti-Spoofing]
|
||||
PRIV[Privacy]
|
||||
@@ -32,7 +33,7 @@ graph TB
|
||||
TRK[BYTETracker]
|
||||
end
|
||||
|
||||
subgraph Indexing
|
||||
subgraph Stores
|
||||
IDX[FAISS Vector Store]
|
||||
end
|
||||
|
||||
@@ -45,6 +46,7 @@ graph TB
|
||||
DET --> LMK
|
||||
DET --> ATTR
|
||||
DET --> GAZE
|
||||
DET --> HPOSE
|
||||
DET --> PARSE
|
||||
DET --> SPOOF
|
||||
DET --> PRIV
|
||||
@@ -119,10 +121,11 @@ uniface/
|
||||
├── attribute/ # Age, gender, emotion, race
|
||||
├── parsing/ # Face semantic segmentation
|
||||
├── gaze/ # Gaze estimation
|
||||
├── headpose/ # Head pose estimation
|
||||
├── spoofing/ # Anti-spoofing
|
||||
├── privacy/ # Face anonymization
|
||||
├── indexing/ # Vector indexing (FAISS)
|
||||
├── types.py # Dataclasses (Face, GazeResult, etc.)
|
||||
├── stores/ # Vector stores (FAISS)
|
||||
├── types.py # Dataclasses (Face, GazeResult, HeadPoseResult, etc.)
|
||||
├── constants.py # Model weights and URLs
|
||||
├── model_store.py # Model download and caching
|
||||
├── onnx_utils.py # ONNX Runtime utilities
|
||||
@@ -158,7 +161,7 @@ for face in faces:
|
||||
embedding = recognizer.get_normalized_embedding(image, face.landmarks)
|
||||
|
||||
# Attributes
|
||||
attrs = age_gender.predict(image, face.bbox)
|
||||
attrs = age_gender.predict(image, face)
|
||||
|
||||
print(f"Face: {attrs.sex}, {attrs.age} years")
|
||||
```
|
||||
@@ -183,8 +186,7 @@ fairface = FairFace()
|
||||
analyzer = FaceAnalyzer(
|
||||
detector,
|
||||
recognizer=recognizer,
|
||||
age_gender=age_gender,
|
||||
fairface=fairface,
|
||||
attributes=[age_gender, fairface],
|
||||
)
|
||||
|
||||
faces = analyzer.analyze(image)
|
||||
|
||||
@@ -201,17 +201,11 @@ For drawing detections, filter by confidence:
|
||||
```python
|
||||
from uniface.draw import draw_detections
|
||||
|
||||
# Only draw high-confidence detections
|
||||
bboxes = [f.bbox for f in faces if f.confidence > 0.7]
|
||||
scores = [f.confidence for f in faces if f.confidence > 0.7]
|
||||
landmarks = [f.landmarks for f in faces if f.confidence > 0.7]
|
||||
|
||||
# Only draw high-confidence detections (confidence ≥ vis_threshold)
|
||||
draw_detections(
|
||||
image=image,
|
||||
bboxes=bboxes,
|
||||
scores=scores,
|
||||
landmarks=landmarks,
|
||||
vis_threshold=0.6 # Additional visualization filter
|
||||
faces=faces,
|
||||
vis_threshold=0.7,
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
@@ -183,6 +183,30 @@ data/
|
||||
|
||||
---
|
||||
|
||||
### Head Pose Estimation
|
||||
|
||||
#### 300W-LP
|
||||
|
||||
Large-scale synthesized face dataset with large pose variations, generated from 300W by face profiling. Used for training head pose estimation models.
|
||||
|
||||
| Property | Value |
|
||||
| ----------- | ----------------------------- |
|
||||
| Images | ~122,000 (synthesized) |
|
||||
| Source | 300W (profiled) |
|
||||
| Pose range | ±90° yaw |
|
||||
| Evaluation | AFLW2000 |
|
||||
| Used by | All HeadPose models |
|
||||
|
||||
!!! info "Download & Reference"
|
||||
**Paper**: [Face Alignment Across Large Poses: A 3D Solution](https://arxiv.org/abs/1511.07212)
|
||||
|
||||
**Training code**: [yakhyo/head-pose-estimation](https://github.com/yakhyo/head-pose-estimation)
|
||||
|
||||
!!! note "UniFace Models"
|
||||
All HeadPose models shipped with UniFace are trained on 300W-LP and evaluated on AFLW2000.
|
||||
|
||||
---
|
||||
|
||||
### Face Parsing
|
||||
|
||||
#### CelebAMask-HQ
|
||||
|
||||
@@ -59,6 +59,11 @@ BiSeNet semantic segmentation with 19 facial component classes.
|
||||
Real-time gaze direction prediction with MobileGaze models.
|
||||
</div>
|
||||
|
||||
<div class="feature-card" markdown>
|
||||
### :material-axis-arrow: Head Pose
|
||||
3D head orientation (pitch, yaw, roll) estimation with 6D rotation models.
|
||||
</div>
|
||||
|
||||
<div class="feature-card" markdown>
|
||||
### :material-motion-play: Tracking
|
||||
Multi-object tracking with BYTETracker for persistent face IDs across video frames.
|
||||
|
||||
@@ -70,16 +70,16 @@ print("Available providers:", ort.get_available_providers())
|
||||
|
||||
---
|
||||
|
||||
### FAISS Vector Indexing
|
||||
### FAISS Vector Store
|
||||
|
||||
For fast multi-identity face search using a FAISS index:
|
||||
For fast multi-identity face search using a FAISS vector store:
|
||||
|
||||
```bash
|
||||
pip install faiss-cpu # CPU
|
||||
pip install faiss-gpu # NVIDIA GPU (CUDA)
|
||||
```
|
||||
|
||||
See the [Indexing module](modules/indexing.md) for usage.
|
||||
See the [Stores module](modules/stores.md) for usage.
|
||||
|
||||
---
|
||||
|
||||
@@ -128,7 +128,7 @@ UniFace has minimal dependencies:
|
||||
|
||||
| Package | Install extra | Purpose |
|
||||
|---------|---------------|---------|
|
||||
| `faiss-cpu` / `faiss-gpu` | `pip install faiss-cpu` | FAISS vector indexing |
|
||||
| `faiss-cpu` / `faiss-gpu` | `pip install faiss-cpu` | FAISS vector store |
|
||||
| `onnxruntime-gpu` | `uniface[gpu]` | CUDA acceleration |
|
||||
| `torch` | `pip install torch` | Emotion model uses TorchScript |
|
||||
| `torchvision` | `pip install torchvision` | Faster NMS for YOLO detectors |
|
||||
|
||||
@@ -257,6 +257,33 @@ Gaze direction prediction models trained on [Gaze360](datasets.md#gaze360) datas
|
||||
|
||||
---
|
||||
|
||||
## Head Pose Estimation Models
|
||||
|
||||
### HeadPose Family
|
||||
|
||||
Head pose estimation models using 6D rotation representation. Trained on [300W-LP](datasets.md#300w-lp) dataset, evaluated on AFLW2000. Returns pitch, yaw, and roll angles in degrees.
|
||||
|
||||
| Model Name | Backbone | Size | MAE* |
|
||||
| -------------- | -------- | ------- | ----- |
|
||||
| `RESNET18` :material-check-circle: | ResNet18 | 43 MB | 5.22° |
|
||||
| `RESNET34` | ResNet34 | 82 MB | 5.07° |
|
||||
| `RESNET50` | ResNet50 | 91 MB | 4.83° |
|
||||
| `MOBILENET_V2` | MobileNetV2 | 9.6 MB | 5.72° |
|
||||
| `MOBILENET_V3_SMALL` | MobileNetV3-Small | 4.8 MB | 6.31° |
|
||||
| `MOBILENET_V3_LARGE` | MobileNetV3-Large | 16 MB | 5.58° |
|
||||
|
||||
*MAE (Mean Absolute Error) in degrees on AFLW2000 test set — lower is better
|
||||
|
||||
!!! info "Training Data"
|
||||
**Dataset**: Trained on [300W-LP](datasets.md#300w-lp) (synthesized large-pose faces from 300W)
|
||||
|
||||
**Method**: 6D rotation representation (rotation matrix → Euler angles)
|
||||
|
||||
!!! note "Input Requirements"
|
||||
Requires face crop as input. Use face detection first to obtain bounding boxes.
|
||||
|
||||
---
|
||||
|
||||
## Face Parsing Models
|
||||
|
||||
### BiSeNet Family
|
||||
@@ -372,6 +399,7 @@ See [Model Cache & Offline Use](concepts/model-cache-offline.md) for full detail
|
||||
- **AdaFace ONNX**: [yakhyo/adaface-onnx](https://github.com/yakhyo/adaface-onnx) - ONNX export and inference
|
||||
- **Face Recognition Training**: [yakhyo/face-recognition](https://github.com/yakhyo/face-recognition) - ArcFace, MobileFace, SphereFace training code
|
||||
- **Gaze Estimation Training**: [yakhyo/gaze-estimation](https://github.com/yakhyo/gaze-estimation) - MobileGaze training code and pretrained weights
|
||||
- **Head Pose Estimation**: [yakhyo/head-pose-estimation](https://github.com/yakhyo/head-pose-estimation) - 6D rotation head pose estimation training and ONNX models
|
||||
- **Face Parsing Training**: [yakhyo/face-parsing](https://github.com/yakhyo/face-parsing) - BiSeNet training code and pretrained weights
|
||||
- **Face Segmentation**: [yakhyo/face-segmentation](https://github.com/yakhyo/face-segmentation) - XSeg ONNX Inference
|
||||
- **Face Anti-Spoofing**: [yakhyo/face-anti-spoofing](https://github.com/yakhyo/face-anti-spoofing) - MiniFASNet ONNX inference (weights from [minivision-ai/Silent-Face-Anti-Spoofing](https://github.com/minivision-ai/Silent-Face-Anti-Spoofing))
|
||||
|
||||
@@ -30,9 +30,10 @@ age_gender = AgeGender()
|
||||
faces = detector.detect(image)
|
||||
|
||||
for face in faces:
|
||||
result = age_gender.predict(image, face.bbox)
|
||||
result = age_gender.predict(image, face)
|
||||
print(f"Gender: {result.sex}") # "Female" or "Male"
|
||||
print(f"Age: {result.age} years")
|
||||
# face.gender and face.age are also set automatically
|
||||
```
|
||||
|
||||
### Output
|
||||
@@ -64,10 +65,11 @@ fairface = FairFace()
|
||||
faces = detector.detect(image)
|
||||
|
||||
for face in faces:
|
||||
result = fairface.predict(image, face.bbox)
|
||||
result = fairface.predict(image, face)
|
||||
print(f"Gender: {result.sex}")
|
||||
print(f"Age Group: {result.age_group}")
|
||||
print(f"Race: {result.race}")
|
||||
# face.gender, face.age_group, face.race are also set automatically
|
||||
```
|
||||
|
||||
### Output
|
||||
@@ -132,7 +134,7 @@ emotion = Emotion(model_name=DDAMFNWeights.AFFECNET7)
|
||||
faces = detector.detect(image)
|
||||
|
||||
for face in faces:
|
||||
result = emotion.predict(image, face.landmarks)
|
||||
result = emotion.predict(image, face)
|
||||
print(f"Emotion: {result.emotion}")
|
||||
print(f"Confidence: {result.confidence:.2%}")
|
||||
```
|
||||
@@ -179,6 +181,22 @@ emotion = Emotion(model_name=DDAMFNWeights.AFFECNET8)
|
||||
|
||||
---
|
||||
|
||||
## Factory Function
|
||||
|
||||
Use `create_attribute_predictor()` for dynamic model selection:
|
||||
|
||||
```python
|
||||
from uniface import create_attribute_predictor
|
||||
|
||||
age_gender = create_attribute_predictor('age_gender')
|
||||
fairface = create_attribute_predictor('fairface')
|
||||
emotion = create_attribute_predictor('emotion')
|
||||
```
|
||||
|
||||
Available model names: `'age_gender'`, `'fairface'`, `'emotion'`.
|
||||
|
||||
---
|
||||
|
||||
## Combining Models
|
||||
|
||||
### Full Attribute Analysis
|
||||
@@ -195,10 +213,10 @@ faces = detector.detect(image)
|
||||
|
||||
for face in faces:
|
||||
# Get exact age from AgeGender
|
||||
ag_result = age_gender.predict(image, face.bbox)
|
||||
ag_result = age_gender.predict(image, face)
|
||||
|
||||
# Get race from FairFace
|
||||
ff_result = fairface.predict(image, face.bbox)
|
||||
ff_result = fairface.predict(image, face)
|
||||
|
||||
print(f"Gender: {ag_result.sex}")
|
||||
print(f"Exact Age: {ag_result.age}")
|
||||
@@ -215,7 +233,7 @@ from uniface.detection import RetinaFace
|
||||
|
||||
analyzer = FaceAnalyzer(
|
||||
RetinaFace(),
|
||||
age_gender=AgeGender(),
|
||||
attributes=[AgeGender()],
|
||||
)
|
||||
|
||||
faces = analyzer.analyze(image)
|
||||
@@ -257,7 +275,7 @@ def draw_attributes(image, face, result):
|
||||
|
||||
# Usage
|
||||
for face in faces:
|
||||
result = age_gender.predict(image, face.bbox)
|
||||
result = age_gender.predict(image, face)
|
||||
image = draw_attributes(image, face, result)
|
||||
|
||||
cv2.imwrite("attributes.jpg", image)
|
||||
|
||||
@@ -264,10 +264,8 @@ from uniface.draw import draw_detections
|
||||
|
||||
draw_detections(
|
||||
image=image,
|
||||
bboxes=[f.bbox for f in faces],
|
||||
scores=[f.confidence for f in faces],
|
||||
landmarks=[f.landmarks for f in faces],
|
||||
vis_threshold=0.6
|
||||
faces=faces,
|
||||
vis_threshold=0.6,
|
||||
)
|
||||
|
||||
cv2.imwrite("result.jpg", image)
|
||||
|
||||
@@ -267,6 +267,7 @@ gaze = create_gaze_estimator() # Returns MobileGaze
|
||||
|
||||
## Next Steps
|
||||
|
||||
- [Head Pose Estimation](headpose.md) - 3D head orientation
|
||||
- [Anti-Spoofing](spoofing.md) - Face liveness detection
|
||||
- [Privacy](privacy.md) - Face anonymization
|
||||
- [Video Recipe](../recipes/video-webcam.md) - Real-time processing
|
||||
|
||||
232
docs/modules/headpose.md
Normal file
232
docs/modules/headpose.md
Normal file
@@ -0,0 +1,232 @@
|
||||
# Head Pose Estimation
|
||||
|
||||
Head pose estimation predicts the 3D orientation of a person's head (pitch, yaw, and roll angles).
|
||||
|
||||
---
|
||||
|
||||
## Available Models
|
||||
|
||||
| Model | Backbone | Size | MAE* |
|
||||
|-------|----------|------|------|
|
||||
| **ResNet18** :material-check-circle: | ResNet18 | 43 MB | 5.22° |
|
||||
| ResNet34 | ResNet34 | 82 MB | 5.07° |
|
||||
| ResNet50 | ResNet50 | 91 MB | 4.83° |
|
||||
| MobileNetV2 | MobileNetV2 | 9.6 MB | 5.72° |
|
||||
| MobileNetV3-Small | MobileNetV3 | 4.8 MB | 6.31° |
|
||||
| MobileNetV3-Large | MobileNetV3 | 16 MB | 5.58° |
|
||||
|
||||
*MAE = Mean Absolute Error on AFLW2000 test set (lower is better)
|
||||
|
||||
---
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.headpose import HeadPose
|
||||
|
||||
detector = RetinaFace()
|
||||
head_pose = HeadPose()
|
||||
|
||||
image = cv2.imread("photo.jpg")
|
||||
faces = detector.detect(image)
|
||||
|
||||
for face in faces:
|
||||
# Crop face
|
||||
x1, y1, x2, y2 = map(int, face.bbox)
|
||||
face_crop = image[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size > 0:
|
||||
# Estimate head pose
|
||||
result = head_pose.estimate(face_crop)
|
||||
print(f"Pitch: {result.pitch:.1f}°, Yaw: {result.yaw:.1f}°, Roll: {result.roll:.1f}°")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Model Variants
|
||||
|
||||
```python
|
||||
from uniface.headpose import HeadPose
|
||||
from uniface.constants import HeadPoseWeights
|
||||
|
||||
# Default (ResNet18, recommended balance of speed and accuracy)
|
||||
hp = HeadPose()
|
||||
|
||||
# Lightweight for mobile/edge
|
||||
hp = HeadPose(model_name=HeadPoseWeights.MOBILENET_V3_SMALL)
|
||||
|
||||
# Higher accuracy
|
||||
hp = HeadPose(model_name=HeadPoseWeights.RESNET50)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Output Format
|
||||
|
||||
```python
|
||||
result = head_pose.estimate(face_crop)
|
||||
|
||||
# HeadPoseResult dataclass
|
||||
result.pitch # Rotation around X-axis in degrees
|
||||
result.yaw # Rotation around Y-axis in degrees
|
||||
result.roll # Rotation around Z-axis in degrees
|
||||
```
|
||||
|
||||
### Angle Convention
|
||||
|
||||
```
|
||||
pitch > 0 (looking down)
|
||||
│
|
||||
│
|
||||
yaw < 0 ─────┼───── yaw > 0
|
||||
(looking left) │ (looking right)
|
||||
│
|
||||
pitch < 0 (looking up)
|
||||
|
||||
roll > 0 = clockwise tilt
|
||||
roll < 0 = counter-clockwise tilt
|
||||
```
|
||||
|
||||
- **Pitch**: Rotation around X-axis (positive = looking down)
|
||||
- **Yaw**: Rotation around Y-axis (positive = looking right)
|
||||
- **Roll**: Rotation around Z-axis (positive = tilting clockwise)
|
||||
|
||||
---
|
||||
|
||||
## Visualization
|
||||
|
||||
### 3D Cube (default)
|
||||
|
||||
The default visualization draws a wireframe cube oriented to match the head pose.
|
||||
|
||||
```python
|
||||
from uniface.draw import draw_head_pose
|
||||
|
||||
faces = detector.detect(image)
|
||||
|
||||
for face in faces:
|
||||
x1, y1, x2, y2 = map(int, face.bbox)
|
||||
face_crop = image[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size > 0:
|
||||
result = head_pose.estimate(face_crop)
|
||||
|
||||
# Draw cube on image (default)
|
||||
draw_head_pose(image, face.bbox, result.pitch, result.yaw, result.roll)
|
||||
|
||||
cv2.imwrite("headpose_output.jpg", image)
|
||||
```
|
||||
|
||||
### Axis Visualization
|
||||
|
||||
```python
|
||||
from uniface.draw import draw_head_pose
|
||||
|
||||
# X/Y/Z coordinate axes
|
||||
draw_head_pose(image, face.bbox, result.pitch, result.yaw, result.roll, draw_type='axis')
|
||||
```
|
||||
|
||||
### Low-Level Drawing Functions
|
||||
|
||||
```python
|
||||
from uniface.draw import draw_head_pose_cube, draw_head_pose_axis
|
||||
|
||||
# Draw cube directly
|
||||
draw_head_pose_cube(image, yaw=10.0, pitch=-5.0, roll=2.0, bbox=[100, 100, 250, 280])
|
||||
|
||||
# Draw axes directly
|
||||
draw_head_pose_axis(image, yaw=10.0, pitch=-5.0, roll=2.0, bbox=[100, 100, 250, 280])
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Real-Time Head Pose Tracking
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.headpose import HeadPose
|
||||
from uniface.draw import draw_head_pose
|
||||
|
||||
detector = RetinaFace()
|
||||
head_pose = HeadPose()
|
||||
|
||||
cap = cv2.VideoCapture(0)
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
for face in faces:
|
||||
x1, y1, x2, y2 = map(int, face.bbox)
|
||||
face_crop = frame[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size > 0:
|
||||
result = head_pose.estimate(face_crop)
|
||||
draw_head_pose(frame, face.bbox, result.pitch, result.yaw, result.roll)
|
||||
|
||||
cv2.imshow("Head Pose Estimation", frame)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Use Cases
|
||||
|
||||
### Driver Drowsiness Detection
|
||||
|
||||
```python
|
||||
def is_head_drooping(result, pitch_threshold=-15):
|
||||
"""Check if the head is drooping (looking down significantly)."""
|
||||
return result.pitch < pitch_threshold
|
||||
|
||||
result = head_pose.estimate(face_crop)
|
||||
if is_head_drooping(result):
|
||||
print("Warning: Head drooping detected")
|
||||
```
|
||||
|
||||
### Attention Monitoring
|
||||
|
||||
```python
|
||||
def is_facing_forward(result, threshold=20):
|
||||
"""Check if the person is facing roughly forward."""
|
||||
return (
|
||||
abs(result.pitch) < threshold
|
||||
and abs(result.yaw) < threshold
|
||||
and abs(result.roll) < threshold
|
||||
)
|
||||
|
||||
result = head_pose.estimate(face_crop)
|
||||
if is_facing_forward(result):
|
||||
print("Facing forward")
|
||||
else:
|
||||
print("Looking away")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Factory Function
|
||||
|
||||
```python
|
||||
from uniface.headpose import create_head_pose_estimator
|
||||
|
||||
hp = create_head_pose_estimator() # Returns HeadPose
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Next Steps
|
||||
|
||||
- [Gaze Estimation](gaze.md) - Eye gaze direction
|
||||
- [Anti-Spoofing](spoofing.md) - Face liveness detection
|
||||
- [Video Recipe](../recipes/video-webcam.md) - Real-time processing
|
||||
@@ -1,4 +1,4 @@
|
||||
# Indexing
|
||||
# Stores
|
||||
|
||||
FAISS-backed vector store for fast similarity search over embeddings.
|
||||
|
||||
@@ -12,7 +12,7 @@ FAISS-backed vector store for fast similarity search over embeddings.
|
||||
## FAISS
|
||||
|
||||
```python
|
||||
from uniface.indexing import FAISS
|
||||
from uniface.stores import FAISS
|
||||
```
|
||||
|
||||
A thin wrapper around a FAISS `IndexFlatIP` (inner-product) index. Vectors
|
||||
@@ -134,7 +134,7 @@ loaded = store.load() # True if files exist
|
||||
import cv2
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.recognition import ArcFace
|
||||
from uniface.indexing import FAISS
|
||||
from uniface.stores import FAISS
|
||||
|
||||
detector = RetinaFace()
|
||||
recognizer = ArcFace()
|
||||
@@ -18,6 +18,8 @@ Run UniFace examples directly in your browser with Google Colab, or download and
|
||||
| [Gaze Estimation](https://github.com/yakhyo/uniface/blob/main/examples/08_gaze_estimation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/08_gaze_estimation.ipynb) | Gaze direction estimation |
|
||||
| [Face Segmentation](https://github.com/yakhyo/uniface/blob/main/examples/09_face_segmentation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/09_face_segmentation.ipynb) | Face segmentation with XSeg |
|
||||
| [Face Vector Store](https://github.com/yakhyo/uniface/blob/main/examples/10_face_vector_store.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/10_face_vector_store.ipynb) | FAISS-backed face database |
|
||||
| [Head Pose Estimation](https://github.com/yakhyo/uniface/blob/main/examples/11_head_pose_estimation.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/11_head_pose_estimation.ipynb) | 3D head orientation estimation |
|
||||
| [Face Recognition](https://github.com/yakhyo/uniface/blob/main/examples/12_face_recognition.ipynb) | [](https://colab.research.google.com/github/yakhyo/uniface/blob/main/examples/12_face_recognition.ipynb) | Standalone face recognition pipeline |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -54,19 +54,8 @@ detector = RetinaFace()
|
||||
image = cv2.imread("photo.jpg")
|
||||
faces = detector.detect(image)
|
||||
|
||||
# Extract visualization data
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
|
||||
# Draw on image
|
||||
draw_detections(
|
||||
image=image,
|
||||
bboxes=bboxes,
|
||||
scores=scores,
|
||||
landmarks=landmarks,
|
||||
vis_threshold=0.6,
|
||||
)
|
||||
draw_detections(image=image, faces=faces, vis_threshold=0.6)
|
||||
|
||||
# Save result
|
||||
cv2.imwrite("output.jpg", image)
|
||||
@@ -80,7 +69,6 @@ Compare two faces:
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import numpy as np
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.recognition import ArcFace
|
||||
|
||||
@@ -97,12 +85,13 @@ faces1 = detector.detect(image1)
|
||||
faces2 = detector.detect(image2)
|
||||
|
||||
if faces1 and faces2:
|
||||
# Extract embeddings
|
||||
# Extract embeddings (normalized 1-D vectors)
|
||||
emb1 = recognizer.get_normalized_embedding(image1, faces1[0].landmarks)
|
||||
emb2 = recognizer.get_normalized_embedding(image2, faces2[0].landmarks)
|
||||
|
||||
# Compute similarity (cosine similarity)
|
||||
similarity = np.dot(emb1, emb2.T)[0][0]
|
||||
# Compute cosine similarity
|
||||
from uniface import compute_similarity
|
||||
similarity = compute_similarity(emb1, emb2, normalized=True)
|
||||
|
||||
# Interpret result
|
||||
if similarity > 0.6:
|
||||
@@ -135,7 +124,7 @@ faces = detector.detect(image)
|
||||
|
||||
# Predict attributes
|
||||
for i, face in enumerate(faces):
|
||||
result = age_gender.predict(image, face.bbox)
|
||||
result = age_gender.predict(image, face)
|
||||
print(f"Face {i+1}: {result.sex}, {result.age} years old")
|
||||
```
|
||||
|
||||
@@ -164,7 +153,7 @@ image = cv2.imread("photo.jpg")
|
||||
faces = detector.detect(image)
|
||||
|
||||
for i, face in enumerate(faces):
|
||||
result = fairface.predict(image, face.bbox)
|
||||
result = fairface.predict(image, face)
|
||||
print(f"Face {i+1}: {result.sex}, {result.age_group}, {result.race}")
|
||||
```
|
||||
|
||||
@@ -234,6 +223,36 @@ cv2.imwrite("gaze_output.jpg", image)
|
||||
|
||||
---
|
||||
|
||||
## Head Pose Estimation
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.headpose import HeadPose
|
||||
from uniface.draw import draw_head_pose
|
||||
|
||||
detector = RetinaFace()
|
||||
head_pose = HeadPose()
|
||||
|
||||
image = cv2.imread("photo.jpg")
|
||||
faces = detector.detect(image)
|
||||
|
||||
for i, face in enumerate(faces):
|
||||
x1, y1, x2, y2 = map(int, face.bbox[:4])
|
||||
face_crop = image[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size > 0:
|
||||
result = head_pose.estimate(face_crop)
|
||||
print(f"Face {i+1}: pitch={result.pitch:.1f}°, yaw={result.yaw:.1f}°, roll={result.roll:.1f}°")
|
||||
|
||||
# Draw 3D cube visualization
|
||||
draw_head_pose(image, face.bbox, result.pitch, result.yaw, result.roll)
|
||||
|
||||
cv2.imwrite("headpose_output.jpg", image)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Face Parsing
|
||||
|
||||
Segment face into semantic components:
|
||||
@@ -342,10 +361,7 @@ while True:
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks)
|
||||
draw_detections(image=frame, faces=faces)
|
||||
|
||||
cv2.imshow("UniFace - Press 'q' to quit", frame)
|
||||
|
||||
@@ -424,6 +440,7 @@ For detailed model comparisons and benchmarks, see the [Model Zoo](models.md).
|
||||
| Recognition | `ArcFace`, `AdaFace`, `MobileFace`, `SphereFace` |
|
||||
| Tracking | `BYTETracker` |
|
||||
| Gaze | `MobileGaze` (ResNet18/34/50, MobileNetV2, MobileOneS0) |
|
||||
| Head Pose | `HeadPose` (ResNet18/34/50, MobileNetV2/V3) |
|
||||
| Parsing | `BiSeNet` (ResNet18/34) |
|
||||
| Attributes | `AgeGender`, `FairFace`, `Emotion` |
|
||||
| Anti-Spoofing | `MiniFASNet` (V1SE, V2) |
|
||||
@@ -470,12 +487,13 @@ from uniface.recognition import ArcFace, AdaFace
|
||||
from uniface.attribute import AgeGender, FairFace
|
||||
from uniface.landmark import Landmark106
|
||||
from uniface.gaze import MobileGaze
|
||||
from uniface.headpose import HeadPose
|
||||
from uniface.parsing import BiSeNet, XSeg
|
||||
from uniface.privacy import BlurFace
|
||||
from uniface.spoofing import MiniFASNet
|
||||
from uniface.tracking import BYTETracker
|
||||
from uniface.analyzer import FaceAnalyzer
|
||||
from uniface.indexing import FAISS # pip install faiss-cpu
|
||||
from uniface.stores import FAISS # pip install faiss-cpu
|
||||
from uniface.draw import draw_detections, draw_tracks
|
||||
```
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ python tools/search.py --reference ref.jpg --source 0 # webcam
|
||||
## Vector Search (FAISS index)
|
||||
|
||||
For identifying faces against a database of many known people, use the
|
||||
[`FAISS`](../modules/indexing.md) vector store.
|
||||
[`FAISS`](../modules/stores.md) vector store.
|
||||
|
||||
!!! info "Install extra"
|
||||
`bash
|
||||
@@ -80,7 +80,7 @@ import cv2
|
||||
from pathlib import Path
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.recognition import ArcFace
|
||||
from uniface.indexing import FAISS
|
||||
from uniface.stores import FAISS
|
||||
|
||||
detector = RetinaFace()
|
||||
recognizer = ArcFace()
|
||||
@@ -112,7 +112,7 @@ python tools/faiss_search.py build --faces-dir dataset/ --db-path ./my_index
|
||||
import cv2
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.recognition import ArcFace
|
||||
from uniface.indexing import FAISS
|
||||
from uniface.stores import FAISS
|
||||
|
||||
detector = RetinaFace()
|
||||
recognizer = ArcFace()
|
||||
@@ -143,7 +143,7 @@ python tools/faiss_search.py run --db-path ./my_index --source 0 # webcam
|
||||
### Manage the index
|
||||
|
||||
```python
|
||||
from uniface.indexing import FAISS
|
||||
from uniface.stores import FAISS
|
||||
|
||||
store = FAISS(db_path="./my_index")
|
||||
store.load()
|
||||
@@ -160,7 +160,7 @@ store.save()
|
||||
|
||||
## See Also
|
||||
|
||||
- [Indexing Module](../modules/indexing.md) - Full `FAISS` API reference
|
||||
- [Stores Module](../modules/stores.md) - Full `FAISS` API reference
|
||||
- [Recognition Module](../modules/recognition.md) - Face recognition details
|
||||
- [Video & Webcam](video-webcam.md) - Real-time processing
|
||||
- [Concepts: Thresholds](../concepts/thresholds-calibration.md) - Tuning similarity thresholds
|
||||
|
||||
@@ -34,7 +34,7 @@ def process_image(image_path):
|
||||
embedding = recognizer.get_normalized_embedding(image, face.landmarks)
|
||||
|
||||
# Step 3: Predict attributes
|
||||
attrs = age_gender.predict(image, face.bbox)
|
||||
attrs = age_gender.predict(image, face)
|
||||
|
||||
results.append({
|
||||
'face_id': i,
|
||||
@@ -48,12 +48,7 @@ def process_image(image_path):
|
||||
print(f" Face {i+1}: {attrs.sex}, {attrs.age} years old")
|
||||
|
||||
# Visualize
|
||||
draw_detections(
|
||||
image=image,
|
||||
bboxes=[f.bbox for f in faces],
|
||||
scores=[f.confidence for f in faces],
|
||||
landmarks=[f.landmarks for f in faces]
|
||||
)
|
||||
draw_detections(image=image, faces=faces)
|
||||
|
||||
return image, results
|
||||
|
||||
@@ -83,7 +78,7 @@ age_gender = AgeGender()
|
||||
analyzer = FaceAnalyzer(
|
||||
detector,
|
||||
recognizer=recognizer,
|
||||
age_gender=age_gender,
|
||||
attributes=[age_gender],
|
||||
)
|
||||
|
||||
# Process image
|
||||
@@ -109,11 +104,12 @@ import numpy as np
|
||||
from uniface.attribute import AgeGender, FairFace
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.gaze import MobileGaze
|
||||
from uniface.headpose import HeadPose
|
||||
from uniface.landmark import Landmark106
|
||||
from uniface.recognition import ArcFace
|
||||
from uniface.parsing import BiSeNet
|
||||
from uniface.spoofing import MiniFASNet
|
||||
from uniface.draw import draw_detections, draw_gaze
|
||||
from uniface.draw import draw_detections, draw_gaze, draw_head_pose
|
||||
|
||||
class FaceAnalysisPipeline:
|
||||
def __init__(self):
|
||||
@@ -124,6 +120,7 @@ class FaceAnalysisPipeline:
|
||||
self.fairface = FairFace()
|
||||
self.landmarker = Landmark106()
|
||||
self.gaze = MobileGaze()
|
||||
self.head_pose = HeadPose()
|
||||
self.parser = BiSeNet()
|
||||
self.spoofer = MiniFASNet()
|
||||
|
||||
@@ -145,12 +142,12 @@ class FaceAnalysisPipeline:
|
||||
)
|
||||
|
||||
# Attributes
|
||||
ag_result = self.age_gender.predict(image, face.bbox)
|
||||
ag_result = self.age_gender.predict(image, face)
|
||||
result['age'] = ag_result.age
|
||||
result['gender'] = ag_result.sex
|
||||
|
||||
# FairFace attributes
|
||||
ff_result = self.fairface.predict(image, face.bbox)
|
||||
ff_result = self.fairface.predict(image, face)
|
||||
result['age_group'] = ff_result.age_group
|
||||
result['race'] = ff_result.race
|
||||
|
||||
@@ -167,6 +164,13 @@ class FaceAnalysisPipeline:
|
||||
result['gaze_pitch'] = gaze_result.pitch
|
||||
result['gaze_yaw'] = gaze_result.yaw
|
||||
|
||||
# Head pose estimation
|
||||
if face_crop.size > 0:
|
||||
hp_result = self.head_pose.estimate(face_crop)
|
||||
result['head_pitch'] = hp_result.pitch
|
||||
result['head_yaw'] = hp_result.yaw
|
||||
result['head_roll'] = hp_result.roll
|
||||
|
||||
# Face parsing
|
||||
if face_crop.size > 0:
|
||||
result['parsing_mask'] = self.parser.parse(face_crop)
|
||||
@@ -189,6 +193,7 @@ for i, r in enumerate(results):
|
||||
print(f" Gender: {r['gender']}, Age: {r['age']}")
|
||||
print(f" Race: {r['race']}, Age Group: {r['age_group']}")
|
||||
print(f" Gaze: pitch={np.degrees(r['gaze_pitch']):.1f}°")
|
||||
print(f" Head Pose: P={r['head_pitch']:.1f}° Y={r['head_yaw']:.1f}° R={r['head_roll']:.1f}°")
|
||||
print(f" Real: {r['is_real']} ({r['spoof_confidence']:.1%})")
|
||||
```
|
||||
|
||||
@@ -220,7 +225,7 @@ def visualize_analysis(image_path, output_path):
|
||||
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||||
|
||||
# Age and gender
|
||||
attrs = age_gender.predict(image, face.bbox)
|
||||
attrs = age_gender.predict(image, face)
|
||||
label = f"{attrs.sex}, {attrs.age}y"
|
||||
cv2.putText(image, label, (x1, y1 - 10),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
||||
@@ -268,6 +273,11 @@ def results_to_json(results):
|
||||
'gaze': {
|
||||
'pitch_deg': float(np.degrees(r['gaze_pitch'])) if 'gaze_pitch' in r else None,
|
||||
'yaw_deg': float(np.degrees(r['gaze_yaw'])) if 'gaze_yaw' in r else None
|
||||
},
|
||||
'head_pose': {
|
||||
'pitch': float(r['head_pitch']) if 'head_pitch' in r else None,
|
||||
'yaw': float(r['head_yaw']) if 'head_yaw' in r else None,
|
||||
'roll': float(r['head_roll']) if 'head_roll' in r else None
|
||||
}
|
||||
}
|
||||
output.append(item)
|
||||
@@ -291,3 +301,4 @@ with open('results.json', 'w') as f:
|
||||
- [Face Search](face-search.md) - Build a search system
|
||||
- [Detection Module](../modules/detection.md) - Detection options
|
||||
- [Recognition Module](../modules/recognition.md) - Recognition details
|
||||
- [Head Pose Module](../modules/headpose.md) - Head orientation estimation
|
||||
|
||||
@@ -26,12 +26,7 @@ while True:
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
draw_detections(
|
||||
image=frame,
|
||||
bboxes=[f.bbox for f in faces],
|
||||
scores=[f.confidence for f in faces],
|
||||
landmarks=[f.landmarks for f in faces]
|
||||
)
|
||||
draw_detections(image=frame, faces=faces)
|
||||
|
||||
cv2.imshow("Face Detection", frame)
|
||||
|
||||
@@ -175,3 +170,4 @@ while True:
|
||||
- [Batch Processing](batch-processing.md) - Process multiple files
|
||||
- [Detection Module](../modules/detection.md) - Detection options
|
||||
- [Gaze Module](../modules/gaze.md) - Gaze estimation
|
||||
- [Head Pose Module](../modules/headpose.md) - Head orientation estimation
|
||||
|
||||
@@ -51,7 +51,7 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"3.0.0\n"
|
||||
"3.3.0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -156,13 +156,8 @@
|
||||
"faces = detector.detect(image)\n",
|
||||
"print(f'Detected {len(faces)} face(s)')\n",
|
||||
"\n",
|
||||
"# Unpack face data for visualization\n",
|
||||
"bboxes = [f.bbox for f in faces]\n",
|
||||
"scores = [f.confidence for f in faces]\n",
|
||||
"landmarks = [f.landmarks for f in faces]\n",
|
||||
"\n",
|
||||
"# Draw detections\n",
|
||||
"draw_detections(image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"draw_detections(image=image, faces=faces, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"\n",
|
||||
"# Display result\n",
|
||||
"output_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
|
||||
@@ -210,11 +205,7 @@
|
||||
"faces = detector.detect(image, max_num=2)\n",
|
||||
"print(f'Detected {len(faces)} face(s)')\n",
|
||||
"\n",
|
||||
"bboxes = [f.bbox for f in faces]\n",
|
||||
"scores = [f.confidence for f in faces]\n",
|
||||
"landmarks = [f.landmarks for f in faces]\n",
|
||||
"\n",
|
||||
"draw_detections(image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"draw_detections(image=image, faces=faces, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"\n",
|
||||
"output_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
|
||||
"display.display(Image.fromarray(output_image))"
|
||||
@@ -257,11 +248,7 @@
|
||||
"faces = detector.detect(image, max_num=5)\n",
|
||||
"print(f'Detected {len(faces)} face(s)')\n",
|
||||
"\n",
|
||||
"bboxes = [f.bbox for f in faces]\n",
|
||||
"scores = [f.confidence for f in faces]\n",
|
||||
"landmarks = [f.landmarks for f in faces]\n",
|
||||
"\n",
|
||||
"draw_detections(image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"draw_detections(image=image, faces=faces, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"\n",
|
||||
"output_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
|
||||
"display.display(Image.fromarray(output_image))"
|
||||
|
||||
@@ -55,7 +55,7 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"3.0.0\n"
|
||||
"3.3.0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -139,10 +139,7 @@
|
||||
"\n",
|
||||
" # Draw detections\n",
|
||||
" bbox_image = image.copy()\n",
|
||||
" bboxes = [f.bbox for f in faces]\n",
|
||||
" scores = [f.confidence for f in faces]\n",
|
||||
" landmarks = [f.landmarks for f in faces]\n",
|
||||
" draw_detections(image=bbox_image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
" draw_detections(image=bbox_image, faces=faces, vis_threshold=0.6, corner_bbox=True)\n",
|
||||
"\n",
|
||||
" # Align first detected face (returns aligned image and inverse transform matrix)\n",
|
||||
" first_landmarks = faces[0].landmarks\n",
|
||||
|
||||
@@ -44,7 +44,7 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"3.0.0\n"
|
||||
"3.3.0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -51,7 +51,7 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"3.0.0\n"
|
||||
"3.3.0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -87,7 +87,7 @@
|
||||
"analyzer = FaceAnalyzer(\n",
|
||||
" detector=RetinaFace(confidence_threshold=0.5),\n",
|
||||
" recognizer=ArcFace(),\n",
|
||||
" age_gender=AgeGender()\n",
|
||||
" attributes=[AgeGender()]\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
@@ -145,10 +145,7 @@
|
||||
"\n",
|
||||
" # Prepare visualization (without text overlay)\n",
|
||||
" vis_image = image.copy()\n",
|
||||
" bboxes = [f.bbox for f in faces]\n",
|
||||
" scores = [f.confidence for f in faces]\n",
|
||||
" landmarks = [f.landmarks for f in faces]\n",
|
||||
" draw_detections(image=vis_image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=0.5, corner_bbox=True)\n",
|
||||
" draw_detections(image=vis_image, faces=faces, vis_threshold=0.5, corner_bbox=True)\n",
|
||||
"\n",
|
||||
" results.append((image_path, cv2.cvtColor(vis_image, cv2.COLOR_BGR2RGB), faces))"
|
||||
]
|
||||
@@ -225,7 +222,7 @@
|
||||
" - Landmarks shape: (5, 2)\n",
|
||||
" - Age: 28 years\n",
|
||||
" - Gender: Female\n",
|
||||
" - Embedding shape: (1, 512)\n",
|
||||
" - Embedding shape: (512,)\n",
|
||||
" - Embedding dimension: 512D\n"
|
||||
]
|
||||
}
|
||||
@@ -243,7 +240,7 @@
|
||||
" print(f' - Age: {face.age} years')\n",
|
||||
" print(f' - Gender: {face.sex}')\n",
|
||||
" print(f' - Embedding shape: {face.embedding.shape}')\n",
|
||||
" print(f' - Embedding dimension: {face.embedding.shape[1]}D')"
|
||||
" print(f' - Embedding dimension: {face.embedding.shape[0]}D')"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -51,7 +51,7 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"UniFace version: 3.0.0\n"
|
||||
"UniFace version: 3.3.0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -53,7 +53,7 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"UniFace version: 3.0.0\n"
|
||||
"UniFace version: 3.3.0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
||||
File diff suppressed because one or more lines are too long
260
examples/11_head_pose_estimation.ipynb
Normal file
260
examples/11_head_pose_estimation.ipynb
Normal file
File diff suppressed because one or more lines are too long
422
examples/12_face_recognition.ipynb
Normal file
422
examples/12_face_recognition.ipynb
Normal file
File diff suppressed because one or more lines are too long
@@ -151,9 +151,10 @@ nav:
|
||||
- Attributes: modules/attributes.md
|
||||
- Parsing: modules/parsing.md
|
||||
- Gaze: modules/gaze.md
|
||||
- Head Pose: modules/headpose.md
|
||||
- Anti-Spoofing: modules/spoofing.md
|
||||
- Privacy: modules/privacy.md
|
||||
- Indexing: modules/indexing.md
|
||||
- Stores: modules/stores.md
|
||||
- Guides:
|
||||
- Overview: concepts/overview.md
|
||||
- Inputs & Outputs: concepts/inputs-outputs.md
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "uniface"
|
||||
version = "3.1.0"
|
||||
version = "3.3.0"
|
||||
description = "UniFace: A Comprehensive Library for Face Detection, Recognition, Tracking, Landmark Analysis, Face Parsing, Gaze Estimation, Age, and Gender Detection"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
@@ -9,7 +9,7 @@ maintainers = [
|
||||
{ name = "Yakhyokhuja Valikhujaev", email = "yakhyo9696@gmail.com" },
|
||||
]
|
||||
|
||||
requires-python = ">=3.10,<3.14"
|
||||
requires-python = ">=3.10,<3.15"
|
||||
keywords = [
|
||||
"face-detection",
|
||||
"face-recognition",
|
||||
@@ -38,13 +38,14 @@ classifiers = [
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"numpy>=1.21.0",
|
||||
"opencv-python>=4.5.0",
|
||||
"onnxruntime>=1.16.0",
|
||||
"scikit-image>=0.19.0",
|
||||
"scikit-image>=0.22.0",
|
||||
"scipy>=1.7.0",
|
||||
"requests>=2.28.0",
|
||||
"tqdm>=4.64.0",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
numpy>=1.21.0
|
||||
opencv-python>=4.5.0
|
||||
onnxruntime>=1.16.0
|
||||
scikit-image>=0.19.0
|
||||
scikit-image>=0.22.0
|
||||
scipy>=1.7.0
|
||||
requests>=2.28.0
|
||||
tqdm>=4.64.0
|
||||
|
||||
@@ -9,6 +9,14 @@ import numpy as np
|
||||
import pytest
|
||||
|
||||
from uniface.attribute import AgeGender, AttributeResult
|
||||
from uniface.types import Face
|
||||
|
||||
|
||||
def _make_face(bbox: list[int] | np.ndarray) -> Face:
|
||||
"""Helper: build a minimal Face from a bounding box."""
|
||||
bbox = np.asarray(bbox)
|
||||
landmarks = np.zeros((5, 2), dtype=np.float32)
|
||||
return Face(bbox=bbox, confidence=0.99, landmarks=landmarks)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -22,30 +30,30 @@ def mock_image():
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_bbox():
|
||||
return [100, 100, 300, 300]
|
||||
def mock_face():
|
||||
return _make_face([100, 100, 300, 300])
|
||||
|
||||
|
||||
def test_model_initialization(age_gender_model):
|
||||
assert age_gender_model is not None, 'AgeGender model initialization failed.'
|
||||
|
||||
|
||||
def test_prediction_output_format(age_gender_model, mock_image, mock_bbox):
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
def test_prediction_output_format(age_gender_model, mock_image, mock_face):
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
assert isinstance(result, AttributeResult), f'Result should be AttributeResult, got {type(result)}'
|
||||
assert isinstance(result.gender, int), f'Gender should be int, got {type(result.gender)}'
|
||||
assert isinstance(result.age, int), f'Age should be int, got {type(result.age)}'
|
||||
assert isinstance(result.sex, str), f'Sex should be str, got {type(result.sex)}'
|
||||
|
||||
|
||||
def test_gender_values(age_gender_model, mock_image, mock_bbox):
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
def test_gender_values(age_gender_model, mock_image, mock_face):
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
assert result.gender in [0, 1], f'Gender should be 0 (Female) or 1 (Male), got {result.gender}'
|
||||
assert result.sex in ['Female', 'Male'], f'Sex should be Female or Male, got {result.sex}'
|
||||
|
||||
|
||||
def test_age_range(age_gender_model, mock_image, mock_bbox):
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
def test_age_range(age_gender_model, mock_image, mock_face):
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
assert 0 <= result.age <= 120, f'Age should be between 0 and 120, got {result.age}'
|
||||
|
||||
|
||||
@@ -57,39 +65,52 @@ def test_different_bbox_sizes(age_gender_model, mock_image):
|
||||
]
|
||||
|
||||
for bbox in test_bboxes:
|
||||
result = age_gender_model.predict(mock_image, bbox)
|
||||
face = _make_face(bbox)
|
||||
result = age_gender_model.predict(mock_image, face)
|
||||
assert result.gender in [0, 1], f'Failed for bbox {bbox}'
|
||||
assert 0 <= result.age <= 120, f'Age out of range for bbox {bbox}'
|
||||
|
||||
|
||||
def test_different_image_sizes(age_gender_model, mock_bbox):
|
||||
def test_different_image_sizes(age_gender_model):
|
||||
test_sizes = [(480, 640, 3), (720, 1280, 3), (1080, 1920, 3)]
|
||||
face = _make_face([100, 100, 300, 300])
|
||||
|
||||
for size in test_sizes:
|
||||
mock_image = np.random.randint(0, 255, size, dtype=np.uint8)
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
result = age_gender_model.predict(mock_image, face)
|
||||
assert result.gender in [0, 1], f'Failed for image size {size}'
|
||||
assert 0 <= result.age <= 120, f'Age out of range for image size {size}'
|
||||
|
||||
|
||||
def test_consistency(age_gender_model, mock_image, mock_bbox):
|
||||
result1 = age_gender_model.predict(mock_image, mock_bbox)
|
||||
result2 = age_gender_model.predict(mock_image, mock_bbox)
|
||||
def test_consistency(age_gender_model, mock_image, mock_face):
|
||||
result1 = age_gender_model.predict(mock_image, mock_face)
|
||||
result2 = age_gender_model.predict(mock_image, mock_face)
|
||||
|
||||
assert result1.gender == result2.gender, 'Same input should produce same gender prediction'
|
||||
assert result1.age == result2.age, 'Same input should produce same age prediction'
|
||||
|
||||
|
||||
def test_face_enrichment(age_gender_model, mock_image, mock_face):
|
||||
"""predict() must write gender & age back to the Face object."""
|
||||
assert mock_face.gender is None
|
||||
assert mock_face.age is None
|
||||
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
|
||||
assert mock_face.gender == result.gender
|
||||
assert mock_face.age == result.age
|
||||
|
||||
|
||||
def test_bbox_list_format(age_gender_model, mock_image):
|
||||
bbox_list = [100, 100, 300, 300]
|
||||
result = age_gender_model.predict(mock_image, bbox_list)
|
||||
face = _make_face([100, 100, 300, 300])
|
||||
result = age_gender_model.predict(mock_image, face)
|
||||
assert result.gender in [0, 1], 'Should work with bbox as list'
|
||||
assert 0 <= result.age <= 120, 'Age should be in valid range'
|
||||
|
||||
|
||||
def test_bbox_array_format(age_gender_model, mock_image):
|
||||
bbox_array = np.array([100, 100, 300, 300])
|
||||
result = age_gender_model.predict(mock_image, bbox_array)
|
||||
face = _make_face(np.array([100, 100, 300, 300]))
|
||||
result = age_gender_model.predict(mock_image, face)
|
||||
assert result.gender in [0, 1], 'Should work with bbox as numpy array'
|
||||
assert 0 <= result.age <= 120, 'Age should be in valid range'
|
||||
|
||||
@@ -103,7 +124,8 @@ def test_multiple_predictions(age_gender_model, mock_image):
|
||||
|
||||
results = []
|
||||
for bbox in bboxes:
|
||||
result = age_gender_model.predict(mock_image, bbox)
|
||||
face = _make_face(bbox)
|
||||
result = age_gender_model.predict(mock_image, face)
|
||||
results.append(result)
|
||||
|
||||
assert len(results) == 3, 'Should have 3 predictions'
|
||||
@@ -112,28 +134,26 @@ def test_multiple_predictions(age_gender_model, mock_image):
|
||||
assert 0 <= result.age <= 120
|
||||
|
||||
|
||||
def test_age_is_positive(age_gender_model, mock_image, mock_bbox):
|
||||
def test_age_is_positive(age_gender_model, mock_image, mock_face):
|
||||
for _ in range(5):
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
assert result.age >= 0, f'Age should be non-negative, got {result.age}'
|
||||
|
||||
|
||||
def test_output_format_for_visualization(age_gender_model, mock_image, mock_bbox):
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
def test_output_format_for_visualization(age_gender_model, mock_image, mock_face):
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
text = f'{result.sex}, {result.age}y'
|
||||
assert isinstance(text, str), 'Should be able to format as string'
|
||||
assert 'Male' in text or 'Female' in text, 'Text should contain gender'
|
||||
assert 'y' in text, "Text should contain 'y' for years"
|
||||
|
||||
|
||||
def test_attribute_result_fields(age_gender_model, mock_image, mock_bbox):
|
||||
def test_attribute_result_fields(age_gender_model, mock_image, mock_face):
|
||||
"""Test that AttributeResult has correct fields for AgeGender model."""
|
||||
result = age_gender_model.predict(mock_image, mock_bbox)
|
||||
result = age_gender_model.predict(mock_image, mock_face)
|
||||
|
||||
# AgeGender should set gender and age
|
||||
assert result.gender is not None
|
||||
assert result.age is not None
|
||||
|
||||
# AgeGender should NOT set race and age_group (FairFace only)
|
||||
assert result.race is None
|
||||
assert result.age_group is None
|
||||
|
||||
@@ -9,12 +9,14 @@ import numpy as np
|
||||
import pytest
|
||||
|
||||
from uniface import (
|
||||
create_attribute_predictor,
|
||||
create_detector,
|
||||
create_landmarker,
|
||||
create_recognizer,
|
||||
list_available_detectors,
|
||||
)
|
||||
from uniface.constants import RetinaFaceWeights, SCRFDWeights
|
||||
from uniface.attribute import AgeGender, FairFace
|
||||
from uniface.constants import AgeGenderWeights, FairFaceWeights, RetinaFaceWeights, SCRFDWeights
|
||||
from uniface.spoofing import MiniFASNet, create_spoofer
|
||||
|
||||
|
||||
@@ -165,7 +167,7 @@ def test_recognizer_inference_from_factory():
|
||||
|
||||
embedding = recognizer.get_embedding(mock_image)
|
||||
assert embedding is not None, 'Recognizer should return embedding'
|
||||
assert embedding.shape[1] == 512, 'Should return 512-dimensional embedding'
|
||||
assert embedding.shape == (1, 512), 'get_embedding should return (1, 512) with batch dimension'
|
||||
|
||||
|
||||
def test_landmarker_inference_from_factory():
|
||||
@@ -236,3 +238,19 @@ def test_create_spoofer_with_providers():
|
||||
"""Test that create_spoofer forwards providers kwarg without TypeError."""
|
||||
spoofer = create_spoofer(providers=['CPUExecutionProvider'])
|
||||
assert isinstance(spoofer, MiniFASNet), 'Should return MiniFASNet instance'
|
||||
|
||||
|
||||
# create_attribute_predictor tests
|
||||
def test_create_attribute_predictor_age_gender():
|
||||
predictor = create_attribute_predictor(AgeGenderWeights.DEFAULT)
|
||||
assert isinstance(predictor, AgeGender), 'Should return AgeGender instance'
|
||||
|
||||
|
||||
def test_create_attribute_predictor_fairface():
|
||||
predictor = create_attribute_predictor(FairFaceWeights.DEFAULT)
|
||||
assert isinstance(predictor, FairFace), 'Should return FairFace instance'
|
||||
|
||||
|
||||
def test_create_attribute_predictor_invalid():
|
||||
with pytest.raises(ValueError, match='Unsupported attribute model'):
|
||||
create_attribute_predictor('invalid_model')
|
||||
|
||||
115
tests/test_headpose.py
Normal file
115
tests/test_headpose.py
Normal file
@@ -0,0 +1,115 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from uniface import HeadPose, HeadPoseResult, create_head_pose_estimator
|
||||
from uniface.headpose import BaseHeadPoseEstimator
|
||||
from uniface.headpose.models import HeadPose as HeadPoseModel
|
||||
|
||||
|
||||
def test_create_head_pose_estimator_default():
|
||||
"""Test creating a head pose estimator with default parameters."""
|
||||
estimator = create_head_pose_estimator()
|
||||
assert isinstance(estimator, HeadPose), 'Should return HeadPose instance'
|
||||
|
||||
|
||||
def test_create_head_pose_estimator_aliases():
|
||||
"""Test that factory accepts all documented aliases."""
|
||||
for alias in ('headpose', 'head_pose', '6drepnet'):
|
||||
estimator = create_head_pose_estimator(alias)
|
||||
assert isinstance(estimator, HeadPose), f"Alias '{alias}' should return HeadPose"
|
||||
|
||||
|
||||
def test_create_head_pose_estimator_invalid():
|
||||
"""Test that invalid method raises ValueError."""
|
||||
with pytest.raises(ValueError, match='Unsupported head pose estimation method'):
|
||||
create_head_pose_estimator('invalid_method')
|
||||
|
||||
|
||||
def test_head_pose_inference():
|
||||
"""Test that HeadPose can run inference on a mock image."""
|
||||
estimator = HeadPose()
|
||||
mock_image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
||||
result = estimator.estimate(mock_image)
|
||||
|
||||
assert isinstance(result, HeadPoseResult), 'Should return HeadPoseResult'
|
||||
assert isinstance(result.pitch, float), 'pitch should be float'
|
||||
assert isinstance(result.yaw, float), 'yaw should be float'
|
||||
assert isinstance(result.roll, float), 'roll should be float'
|
||||
|
||||
|
||||
def test_head_pose_callable():
|
||||
"""Test that HeadPose is callable via __call__."""
|
||||
estimator = HeadPose()
|
||||
mock_image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
||||
result = estimator(mock_image)
|
||||
|
||||
assert isinstance(result, HeadPoseResult), '__call__ should return HeadPoseResult'
|
||||
|
||||
|
||||
def test_head_pose_result_repr():
|
||||
"""Test HeadPoseResult repr formatting."""
|
||||
result = HeadPoseResult(pitch=10.5, yaw=-20.3, roll=5.1)
|
||||
repr_str = repr(result)
|
||||
assert 'HeadPoseResult' in repr_str
|
||||
assert '10.5' in repr_str
|
||||
assert '-20.3' in repr_str
|
||||
assert '5.1' in repr_str
|
||||
|
||||
|
||||
def test_head_pose_result_frozen():
|
||||
"""Test that HeadPoseResult is immutable."""
|
||||
result = HeadPoseResult(pitch=1.0, yaw=2.0, roll=3.0)
|
||||
with pytest.raises(AttributeError):
|
||||
result.pitch = 99.0 # type: ignore[misc]
|
||||
|
||||
|
||||
def test_rotation_matrix_to_euler_identity():
|
||||
"""Test that identity rotation matrix gives zero angles."""
|
||||
identity = np.eye(3).reshape(1, 3, 3)
|
||||
euler = HeadPoseModel.rotation_matrix_to_euler(identity)
|
||||
|
||||
assert euler.shape == (1, 3), 'Should return (1, 3) shaped array'
|
||||
np.testing.assert_allclose(euler[0], [0.0, 0.0, 0.0], atol=1e-5)
|
||||
|
||||
|
||||
def test_rotation_matrix_to_euler_90deg_yaw():
|
||||
"""Test 90-degree yaw rotation."""
|
||||
angle = np.radians(90)
|
||||
R = np.array(
|
||||
[
|
||||
[np.cos(angle), 0, np.sin(angle)],
|
||||
[0, 1, 0],
|
||||
[-np.sin(angle), 0, np.cos(angle)],
|
||||
]
|
||||
).reshape(1, 3, 3)
|
||||
euler = HeadPoseModel.rotation_matrix_to_euler(R)
|
||||
|
||||
np.testing.assert_allclose(euler[0, 1], 90.0, atol=1e-3)
|
||||
|
||||
|
||||
def test_rotation_matrix_to_euler_batch():
|
||||
"""Test batch processing of rotation matrices."""
|
||||
batch = np.stack([np.eye(3), np.eye(3), np.eye(3)], axis=0)
|
||||
euler = HeadPoseModel.rotation_matrix_to_euler(batch)
|
||||
|
||||
assert euler.shape == (3, 3), 'Batch of 3 should return (3, 3)'
|
||||
np.testing.assert_allclose(euler, 0.0, atol=1e-5)
|
||||
|
||||
|
||||
def test_factory_returns_correct_type():
|
||||
"""Test that factory function returns BaseHeadPoseEstimator subclass."""
|
||||
estimator = create_head_pose_estimator()
|
||||
assert isinstance(estimator, BaseHeadPoseEstimator), 'Should be BaseHeadPoseEstimator subclass'
|
||||
|
||||
|
||||
def test_head_pose_with_providers():
|
||||
"""Test that HeadPose accepts providers kwarg."""
|
||||
estimator = HeadPose(providers=['CPUExecutionProvider'])
|
||||
assert isinstance(estimator, HeadPose), 'Should create with explicit providers'
|
||||
@@ -74,7 +74,7 @@ def test_arcface_embedding_shape(arcface_model, mock_aligned_face):
|
||||
"""
|
||||
embedding = arcface_model.get_embedding(mock_aligned_face)
|
||||
|
||||
# ArcFace typically produces 512-dimensional embeddings
|
||||
# ArcFace get_embedding returns raw ONNX output with batch dimension
|
||||
assert embedding.shape[1] == 512, f'Expected 512-dim embedding, got {embedding.shape[1]}'
|
||||
assert embedding.shape[0] == 1, 'Embedding should have batch dimension of 1'
|
||||
|
||||
@@ -88,7 +88,8 @@ def test_arcface_normalized_embedding(arcface_model, mock_landmarks):
|
||||
|
||||
embedding = arcface_model.get_normalized_embedding(mock_image, mock_landmarks)
|
||||
|
||||
# Check that embedding is normalized (L2 norm ≈ 1.0)
|
||||
# Check shape and normalization
|
||||
assert embedding.shape == (512,), f'Expected shape (512,), got {embedding.shape}'
|
||||
norm = np.linalg.norm(embedding)
|
||||
assert np.isclose(norm, 1.0, atol=1e-5), f'Normalized embedding should have norm 1.0, got {norm}'
|
||||
|
||||
@@ -125,7 +126,7 @@ def test_mobileface_embedding_shape(mobileface_model, mock_aligned_face):
|
||||
"""
|
||||
embedding = mobileface_model.get_embedding(mock_aligned_face)
|
||||
|
||||
# MobileFace typically produces 512-dimensional embeddings
|
||||
# MobileFace get_embedding returns raw ONNX output with batch dimension
|
||||
assert embedding.shape[1] == 512, f'Expected 512-dim embedding, got {embedding.shape[1]}'
|
||||
assert embedding.shape[0] == 1, 'Embedding should have batch dimension of 1'
|
||||
|
||||
@@ -138,6 +139,7 @@ def test_mobileface_normalized_embedding(mobileface_model, mock_landmarks):
|
||||
|
||||
embedding = mobileface_model.get_normalized_embedding(mock_image, mock_landmarks)
|
||||
|
||||
assert embedding.shape == (512,), f'Expected shape (512,), got {embedding.shape}'
|
||||
norm = np.linalg.norm(embedding)
|
||||
assert np.isclose(norm, 1.0, atol=1e-5), f'Normalized embedding should have norm 1.0, got {norm}'
|
||||
|
||||
@@ -156,7 +158,7 @@ def test_sphereface_embedding_shape(sphereface_model, mock_aligned_face):
|
||||
"""
|
||||
embedding = sphereface_model.get_embedding(mock_aligned_face)
|
||||
|
||||
# SphereFace typically produces 512-dimensional embeddings
|
||||
# SphereFace get_embedding returns raw ONNX output with batch dimension
|
||||
assert embedding.shape[1] == 512, f'Expected 512-dim embedding, got {embedding.shape[1]}'
|
||||
assert embedding.shape[0] == 1, 'Embedding should have batch dimension of 1'
|
||||
|
||||
@@ -169,6 +171,7 @@ def test_sphereface_normalized_embedding(sphereface_model, mock_landmarks):
|
||||
|
||||
embedding = sphereface_model.get_normalized_embedding(mock_image, mock_landmarks)
|
||||
|
||||
assert embedding.shape == (512,), f'Expected shape (512,), got {embedding.shape}'
|
||||
norm = np.linalg.norm(embedding)
|
||||
assert np.isclose(norm, 1.0, atol=1e-5), f'Normalized embedding should have norm 1.0, got {norm}'
|
||||
|
||||
|
||||
@@ -12,9 +12,11 @@ CLI utilities for testing and running UniFace features.
|
||||
| `anonymize.py` | Face anonymization/blurring for privacy |
|
||||
| `emotion.py` | Emotion detection (7 or 8 emotions) |
|
||||
| `gaze.py` | Gaze direction estimation |
|
||||
| `headpose.py` | Head pose estimation (pitch, yaw, roll) |
|
||||
| `landmarks.py` | 106-point facial landmark detection |
|
||||
| `recognize.py` | Face embedding extraction and comparison |
|
||||
| `search.py` | Real-time face matching against reference |
|
||||
| `faiss_search.py` | FAISS index build and multi-identity face search |
|
||||
| `fairface.py` | FairFace attribute prediction (race, gender, age) |
|
||||
| `attribute.py` | Age and gender prediction |
|
||||
| `spoofing.py` | Face anti-spoofing detection |
|
||||
@@ -61,6 +63,11 @@ python tools/emotion.py --source 0
|
||||
python tools/gaze.py --source assets/test.jpg
|
||||
python tools/gaze.py --source 0
|
||||
|
||||
# Head pose estimation
|
||||
python tools/headpose.py --source assets/test.jpg
|
||||
python tools/headpose.py --source 0
|
||||
python tools/headpose.py --source 0 --draw-type axis
|
||||
|
||||
# Landmarks
|
||||
python tools/landmarks.py --source assets/test.jpg
|
||||
python tools/landmarks.py --source 0
|
||||
@@ -108,7 +115,7 @@ python tools/download_model.py # downloads all
|
||||
| Option | Description |
|
||||
|--------|-------------|
|
||||
| `--source` | Input source: image/video path or camera ID (0, 1, ...) |
|
||||
| `--detector` | Choose detector: `retinaface`, `scrfd`, `yolov5face` |
|
||||
| `--detector` | Choose detector: `retinaface`, `scrfd`, `yolov5face`, `yolov8face` |
|
||||
| `--threshold` | Visualization confidence threshold (default: varies) |
|
||||
| `--save-dir` | Output directory (default: `outputs`) |
|
||||
|
||||
|
||||
@@ -27,12 +27,17 @@ from uniface.draw import draw_detections
|
||||
from uniface.recognition import ArcFace
|
||||
|
||||
|
||||
def draw_face_info(image, face, face_id):
|
||||
"""Draw face ID and attributes above bounding box."""
|
||||
def draw_face_info(image, face):
|
||||
"""Draw face attributes above bounding box."""
|
||||
x1, y1, _x2, y2 = map(int, face.bbox)
|
||||
lines = [f'ID: {face_id}', f'Conf: {face.confidence:.2f}']
|
||||
if face.age and face.sex:
|
||||
lines = []
|
||||
if face.age is not None and face.sex is not None:
|
||||
lines.append(f'{face.sex}, {face.age}y')
|
||||
if face.emotion is not None:
|
||||
lines.append(face.emotion)
|
||||
|
||||
if not lines:
|
||||
return
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
y_pos = y1 - 10 - (len(lines) - 1 - i) * 25
|
||||
@@ -95,13 +100,10 @@ def process_image(analyzer, image_path: str, save_dir: str = 'outputs', show_sim
|
||||
status = 'Same' if sim > 0.4 else 'Different'
|
||||
print(f' Face {i + 1} ↔ Face {j + 1}: {sim:.3f} ({status})')
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, corner_bbox=True)
|
||||
draw_detections(image=image, faces=faces, corner_bbox=True)
|
||||
|
||||
for i, face in enumerate(faces, 1):
|
||||
draw_face_info(image, face, i)
|
||||
for face in faces:
|
||||
draw_face_info(image, face)
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
output_path = os.path.join(save_dir, f'{Path(image_path).stem}_analysis.jpg')
|
||||
@@ -137,13 +139,10 @@ def process_video(analyzer, video_path: str, save_dir: str = 'outputs'):
|
||||
frame_count += 1
|
||||
faces = analyzer.analyze(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, corner_bbox=True)
|
||||
draw_detections(image=frame, faces=faces, corner_bbox=True)
|
||||
|
||||
for i, face in enumerate(faces, 1):
|
||||
draw_face_info(frame, face, i)
|
||||
for face in faces:
|
||||
draw_face_info(frame, face)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
out.write(frame)
|
||||
@@ -167,19 +166,16 @@ def run_camera(analyzer, camera_id: int = 0):
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = analyzer.analyze(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, corner_bbox=True)
|
||||
draw_detections(image=frame, faces=faces, corner_bbox=True)
|
||||
|
||||
for i, face in enumerate(faces, 1):
|
||||
draw_face_info(frame, face, i)
|
||||
for face in faces:
|
||||
draw_face_info(frame, face)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
cv2.imshow('Face Analyzer', frame)
|
||||
@@ -201,7 +197,7 @@ def main():
|
||||
detector = RetinaFace()
|
||||
recognizer = ArcFace()
|
||||
age_gender = AgeGender()
|
||||
analyzer = FaceAnalyzer(detector, recognizer, age_gender)
|
||||
analyzer = FaceAnalyzer(detector, recognizer=recognizer, attributes=[age_gender])
|
||||
|
||||
source_type = get_source_type(args.source)
|
||||
|
||||
|
||||
@@ -43,10 +43,7 @@ def process_image(
|
||||
from uniface.draw import draw_detections
|
||||
|
||||
preview = image.copy()
|
||||
bboxes = [face.bbox for face in faces]
|
||||
scores = [face.confidence for face in faces]
|
||||
landmarks = [face.landmarks for face in faces]
|
||||
draw_detections(preview, bboxes, scores, landmarks)
|
||||
draw_detections(image=preview, faces=faces)
|
||||
|
||||
cv2.imshow('Detections (Press any key to continue)', preview)
|
||||
cv2.waitKey(0)
|
||||
@@ -121,9 +118,9 @@ def run_camera(detector, blurrer: BlurFace, camera_id: int = 0):
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = detector.detect(frame)
|
||||
if faces:
|
||||
|
||||
@@ -52,15 +52,10 @@ def process_image(
|
||||
if not faces:
|
||||
return
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=image, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for i, face in enumerate(faces):
|
||||
result = age_gender.predict(image, face.bbox)
|
||||
result = age_gender.predict(image, face)
|
||||
print(f' Face {i + 1}: {result.sex}, {result.age} years old')
|
||||
draw_age_gender_label(image, face.bbox, result.sex, result.age)
|
||||
|
||||
@@ -104,15 +99,10 @@ def process_video(
|
||||
frame_count += 1
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=frame, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for face in faces:
|
||||
result = age_gender.predict(frame, face.bbox)
|
||||
result = age_gender.predict(frame, face)
|
||||
draw_age_gender_label(frame, face.bbox, result.sex, result.age)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
@@ -137,21 +127,16 @@ def run_camera(detector, age_gender, camera_id: int = 0, threshold: float = 0.6)
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=frame, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for face in faces:
|
||||
result = age_gender.predict(frame, face.bbox)
|
||||
result = age_gender.predict(frame, face)
|
||||
draw_age_gender_label(frame, face.bbox, result.sex, result.age)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
|
||||
@@ -34,13 +34,7 @@ def process_image(detector, image_path: Path, output_path: Path, threshold: floa
|
||||
|
||||
faces = detector.detect(image)
|
||||
|
||||
# unpack face data for visualization
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=image, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
cv2.putText(
|
||||
image,
|
||||
|
||||
@@ -35,10 +35,7 @@ def process_image(detector, image_path: str, threshold: float = 0.6, save_dir: s
|
||||
faces = detector.detect(image)
|
||||
|
||||
if faces:
|
||||
bboxes = [face.bbox for face in faces]
|
||||
scores = [face.confidence for face in faces]
|
||||
landmarks = [face.landmarks for face in faces]
|
||||
draw_detections(image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold)
|
||||
draw_detections(image=image, faces=faces, vis_threshold=threshold)
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
output_path = os.path.join(save_dir, f'{os.path.splitext(os.path.basename(image_path))[0]}_out.jpg')
|
||||
@@ -89,14 +86,9 @@ def process_video(
|
||||
faces = detector.detect(frame)
|
||||
total_faces += len(faces)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame,
|
||||
bboxes=bboxes,
|
||||
scores=scores,
|
||||
landmarks=landmarks,
|
||||
faces=faces,
|
||||
vis_threshold=threshold,
|
||||
draw_score=True,
|
||||
corner_bbox=True,
|
||||
@@ -135,20 +127,15 @@ def run_camera(detector, camera_id: int = 0, threshold: float = 0.6):
|
||||
prev_time = time.perf_counter()
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame,
|
||||
bboxes=bboxes,
|
||||
scores=scores,
|
||||
landmarks=landmarks,
|
||||
faces=faces,
|
||||
vis_threshold=threshold,
|
||||
draw_score=True,
|
||||
corner_bbox=True,
|
||||
|
||||
@@ -4,6 +4,7 @@ from uniface.constants import (
|
||||
AgeGenderWeights,
|
||||
ArcFaceWeights,
|
||||
DDAMFNWeights,
|
||||
HeadPoseWeights,
|
||||
LandmarkWeights,
|
||||
MobileFaceWeights,
|
||||
RetinaFaceWeights,
|
||||
@@ -21,6 +22,7 @@ MODEL_TYPES = {
|
||||
'ddamfn': DDAMFNWeights,
|
||||
'agegender': AgeGenderWeights,
|
||||
'landmark': LandmarkWeights,
|
||||
'headpose': HeadPoseWeights,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -52,15 +52,10 @@ def process_image(
|
||||
if not faces:
|
||||
return
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=image, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for i, face in enumerate(faces):
|
||||
result = emotion_predictor.predict(image, face.landmarks)
|
||||
result = emotion_predictor.predict(image, face)
|
||||
print(f' Face {i + 1}: {result.emotion} (confidence: {result.confidence:.3f})')
|
||||
draw_emotion_label(image, face.bbox, result.emotion, result.confidence)
|
||||
|
||||
@@ -104,15 +99,10 @@ def process_video(
|
||||
frame_count += 1
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=frame, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for face in faces:
|
||||
result = emotion_predictor.predict(frame, face.landmarks)
|
||||
result = emotion_predictor.predict(frame, face)
|
||||
draw_emotion_label(frame, face.bbox, result.emotion, result.confidence)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
@@ -137,21 +127,16 @@ def run_camera(detector, emotion_predictor, camera_id: int = 0, threshold: float
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=frame, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for face in faces:
|
||||
result = emotion_predictor.predict(frame, face.landmarks)
|
||||
result = emotion_predictor.predict(frame, face)
|
||||
draw_emotion_label(frame, face.bbox, result.emotion, result.confidence)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
|
||||
@@ -52,15 +52,10 @@ def process_image(
|
||||
if not faces:
|
||||
return
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=image, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=image, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for i, face in enumerate(faces):
|
||||
result = fairface.predict(image, face.bbox)
|
||||
result = fairface.predict(image, face)
|
||||
print(f' Face {i + 1}: {result.sex}, {result.age_group}, {result.race}')
|
||||
draw_fairface_label(image, face.bbox, result.sex, result.age_group, result.race)
|
||||
|
||||
@@ -104,15 +99,10 @@ def process_video(
|
||||
frame_count += 1
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=frame, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for face in faces:
|
||||
result = fairface.predict(frame, face.bbox)
|
||||
result = fairface.predict(frame, face)
|
||||
draw_fairface_label(frame, face.bbox, result.sex, result.age_group, result.race)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
@@ -137,21 +127,16 @@ def run_camera(detector, fairface, camera_id: int = 0, threshold: float = 0.6):
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
draw_detections(
|
||||
image=frame, bboxes=bboxes, scores=scores, landmarks=landmarks, vis_threshold=threshold, corner_bbox=True
|
||||
)
|
||||
draw_detections(image=frame, faces=faces, vis_threshold=threshold, corner_bbox=True)
|
||||
|
||||
for face in faces:
|
||||
result = fairface.predict(frame, face.bbox)
|
||||
result = fairface.predict(frame, face)
|
||||
draw_fairface_label(frame, face.bbox, result.sex, result.age_group, result.race)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
|
||||
@@ -24,7 +24,7 @@ import cv2
|
||||
|
||||
from uniface import create_detector, create_recognizer
|
||||
from uniface.draw import draw_corner_bbox, draw_text_label
|
||||
from uniface.indexing import FAISS
|
||||
from uniface.stores import FAISS
|
||||
|
||||
|
||||
def _draw_face(image, bbox, text: str, color: tuple[int, int, int]) -> None:
|
||||
@@ -97,9 +97,9 @@ def run_camera(detector, recognizer, store: FAISS, camera_id: int = 0, threshold
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
frame = process_frame(frame, detector, recognizer, store, threshold)
|
||||
|
||||
|
||||
181
tools/headpose.py
Normal file
181
tools/headpose.py
Normal file
@@ -0,0 +1,181 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
"""Head pose estimation on detected faces.
|
||||
|
||||
Usage:
|
||||
python tools/headpose.py --source path/to/image.jpg
|
||||
python tools/headpose.py --source path/to/video.mp4
|
||||
python tools/headpose.py --source 0 # webcam
|
||||
python tools/headpose.py --source path/to/image.jpg --draw-type axis
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from _common import get_source_type
|
||||
import cv2
|
||||
|
||||
from uniface.detection import RetinaFace
|
||||
from uniface.draw import draw_head_pose
|
||||
from uniface.headpose import HeadPose
|
||||
|
||||
|
||||
def process_image(detector, head_pose_estimator, image_path: str, save_dir: str = 'outputs', draw_type: str = 'cube'):
|
||||
"""Process a single image."""
|
||||
image = cv2.imread(image_path)
|
||||
if image is None:
|
||||
print(f"Error: Failed to load image from '{image_path}'")
|
||||
return
|
||||
|
||||
faces = detector.detect(image)
|
||||
print(f'Detected {len(faces)} face(s)')
|
||||
|
||||
for i, face in enumerate(faces):
|
||||
bbox = face.bbox
|
||||
x1, y1, x2, y2 = map(int, bbox[:4])
|
||||
face_crop = image[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size == 0:
|
||||
continue
|
||||
|
||||
result = head_pose_estimator.estimate(face_crop)
|
||||
print(f' Face {i + 1}: pitch={result.pitch:.1f}°, yaw={result.yaw:.1f}°, roll={result.roll:.1f}°')
|
||||
|
||||
draw_head_pose(image, bbox, result.pitch, result.yaw, result.roll, draw_type=draw_type)
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
output_path = os.path.join(save_dir, f'{Path(image_path).stem}_headpose.jpg')
|
||||
cv2.imwrite(output_path, image)
|
||||
print(f'Output saved: {output_path}')
|
||||
|
||||
|
||||
def process_video(detector, head_pose_estimator, video_path: str, save_dir: str = 'outputs', draw_type: str = 'cube'):
|
||||
"""Process a video file."""
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
print(f"Error: Cannot open video file '{video_path}'")
|
||||
return
|
||||
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
output_path = os.path.join(save_dir, f'{Path(video_path).stem}_headpose.mp4')
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
|
||||
|
||||
print(f'Processing video: {video_path} ({total_frames} frames)')
|
||||
frame_count = 0
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
frame_count += 1
|
||||
faces = detector.detect(frame)
|
||||
|
||||
for face in faces:
|
||||
bbox = face.bbox
|
||||
x1, y1, x2, y2 = map(int, bbox[:4])
|
||||
face_crop = frame[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size == 0:
|
||||
continue
|
||||
|
||||
result = head_pose_estimator.estimate(face_crop)
|
||||
draw_head_pose(frame, bbox, result.pitch, result.yaw, result.roll, draw_type=draw_type)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
out.write(frame)
|
||||
|
||||
if frame_count % 100 == 0:
|
||||
print(f' Processed {frame_count}/{total_frames} frames...')
|
||||
|
||||
cap.release()
|
||||
out.release()
|
||||
print(f'Done! Output saved: {output_path}')
|
||||
|
||||
|
||||
def run_camera(detector, head_pose_estimator, camera_id: int = 0, draw_type: str = 'cube'):
|
||||
"""Run real-time detection on webcam."""
|
||||
cap = cv2.VideoCapture(camera_id)
|
||||
if not cap.isOpened():
|
||||
print(f'Cannot open camera {camera_id}')
|
||||
return
|
||||
|
||||
print("Press 'q' to quit")
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
frame = cv2.flip(frame, 1)
|
||||
faces = detector.detect(frame)
|
||||
|
||||
for face in faces:
|
||||
bbox = face.bbox
|
||||
x1, y1, x2, y2 = map(int, bbox[:4])
|
||||
face_crop = frame[y1:y2, x1:x2]
|
||||
|
||||
if face_crop.size == 0:
|
||||
continue
|
||||
|
||||
result = head_pose_estimator.estimate(face_crop)
|
||||
draw_head_pose(frame, bbox, result.pitch, result.yaw, result.roll, draw_type=draw_type)
|
||||
|
||||
cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
||||
cv2.imshow('Head Pose Estimation', frame)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Run head pose estimation')
|
||||
parser.add_argument('--source', type=str, required=True, help='Image/video path or camera ID (0, 1, ...)')
|
||||
parser.add_argument('--save-dir', type=str, default='outputs', help='Output directory')
|
||||
parser.add_argument(
|
||||
'--draw-type',
|
||||
type=str,
|
||||
default='cube',
|
||||
choices=['cube', 'axis'],
|
||||
help='Visualization type: cube (default) or axis',
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
detector = RetinaFace()
|
||||
head_pose_estimator = HeadPose()
|
||||
|
||||
source_type = get_source_type(args.source)
|
||||
|
||||
if source_type == 'camera':
|
||||
run_camera(detector, head_pose_estimator, int(args.source), args.draw_type)
|
||||
elif source_type == 'image':
|
||||
if not os.path.exists(args.source):
|
||||
print(f'Error: Image not found: {args.source}')
|
||||
return
|
||||
process_image(detector, head_pose_estimator, args.source, args.save_dir, args.draw_type)
|
||||
elif source_type == 'video':
|
||||
if not os.path.exists(args.source):
|
||||
print(f'Error: Video not found: {args.source}')
|
||||
return
|
||||
process_video(detector, head_pose_estimator, args.source, args.save_dir, args.draw_type)
|
||||
else:
|
||||
print(f"Error: Unknown source type for '{args.source}'")
|
||||
print('Supported formats: images (.jpg, .png, ...), videos (.mp4, .avi, ...), or camera ID (0, 1, ...)')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -114,9 +114,9 @@ def run_camera(detector, landmarker, camera_id: int = 0):
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
faces = detector.detect(frame)
|
||||
|
||||
|
||||
@@ -41,12 +41,13 @@ def run_inference(detector, recognizer, image_path: str):
|
||||
|
||||
print(f'Detected {len(faces)} face(s). Extracting embedding for the first face...')
|
||||
|
||||
landmarks = faces[0].landmarks # 5-point landmarks for alignment (already np.ndarray)
|
||||
landmarks = faces[0].landmarks
|
||||
embedding = recognizer.get_embedding(image, landmarks)
|
||||
norm_embedding = recognizer.get_normalized_embedding(image, landmarks) # L2 normalized
|
||||
raw_norm = np.linalg.norm(embedding)
|
||||
norm_embedding = embedding.ravel() / raw_norm if raw_norm > 0 else embedding.ravel()
|
||||
|
||||
print(f' Embedding shape: {embedding.shape}')
|
||||
print(f' L2 norm (raw): {np.linalg.norm(embedding):.4f}')
|
||||
print(f' L2 norm (raw): {raw_norm:.4f}')
|
||||
print(f' L2 norm (normalized): {np.linalg.norm(norm_embedding):.4f}')
|
||||
|
||||
|
||||
|
||||
@@ -109,9 +109,9 @@ def run_camera(detector, recognizer, ref_embedding: np.ndarray, camera_id: int =
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
frame = process_frame(frame, detector, recognizer, ref_embedding, threshold)
|
||||
|
||||
|
||||
@@ -134,9 +134,9 @@ def run_camera(
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
frame = cv2.flip(frame, 1)
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.flip(frame, 1)
|
||||
|
||||
# Detect faces
|
||||
faces = detector.detect(frame)
|
||||
|
||||
@@ -20,6 +20,7 @@ This library provides unified APIs for:
|
||||
- Facial landmarks (106-point detection)
|
||||
- Face parsing (semantic segmentation)
|
||||
- Gaze estimation
|
||||
- Head pose estimation
|
||||
- Age, gender, and emotion prediction
|
||||
- Face anti-spoofing
|
||||
- Privacy/anonymization
|
||||
@@ -29,7 +30,7 @@ from __future__ import annotations
|
||||
|
||||
__license__ = 'MIT'
|
||||
__author__ = 'Yakhyokhuja Valikhujaev'
|
||||
__version__ = '3.1.0'
|
||||
__version__ = '3.3.0'
|
||||
|
||||
import contextlib
|
||||
|
||||
@@ -38,7 +39,7 @@ from uniface.log import Logger, enable_logging
|
||||
from uniface.model_store import download_models, get_cache_dir, set_cache_dir, verify_model_weights
|
||||
|
||||
from .analyzer import FaceAnalyzer
|
||||
from .attribute import AgeGender, Emotion, FairFace
|
||||
from .attribute import AgeGender, Emotion, FairFace, create_attribute_predictor
|
||||
from .detection import (
|
||||
SCRFD,
|
||||
RetinaFace,
|
||||
@@ -48,17 +49,18 @@ from .detection import (
|
||||
list_available_detectors,
|
||||
)
|
||||
from .gaze import MobileGaze, create_gaze_estimator
|
||||
from .headpose import HeadPose, create_head_pose_estimator
|
||||
from .landmark import Landmark106, create_landmarker
|
||||
from .parsing import BiSeNet, XSeg, create_face_parser
|
||||
from .privacy import BlurFace
|
||||
from .recognition import AdaFace, ArcFace, MobileFace, SphereFace, create_recognizer
|
||||
from .spoofing import MiniFASNet, create_spoofer
|
||||
from .tracking import BYTETracker
|
||||
from .types import AttributeResult, EmotionResult, Face, GazeResult, SpoofingResult
|
||||
from .types import AttributeResult, EmotionResult, Face, GazeResult, HeadPoseResult, SpoofingResult
|
||||
|
||||
# Optional: FAISS vector store (requires `pip install faiss-cpu`)
|
||||
with contextlib.suppress(ImportError):
|
||||
from .indexing import FAISS
|
||||
from .stores import FAISS
|
||||
|
||||
__all__ = [
|
||||
# Metadata
|
||||
@@ -72,6 +74,7 @@ __all__ = [
|
||||
'create_detector',
|
||||
'create_face_parser',
|
||||
'create_gaze_estimator',
|
||||
'create_head_pose_estimator',
|
||||
'create_landmarker',
|
||||
'create_recognizer',
|
||||
'create_spoofer',
|
||||
@@ -91,12 +94,16 @@ __all__ = [
|
||||
# Gaze models
|
||||
'GazeResult',
|
||||
'MobileGaze',
|
||||
# Head pose models
|
||||
'HeadPose',
|
||||
'HeadPoseResult',
|
||||
# Parsing models
|
||||
'BiSeNet',
|
||||
'XSeg',
|
||||
# Attribute models
|
||||
'AgeGender',
|
||||
'AttributeResult',
|
||||
'create_attribute_predictor',
|
||||
'Emotion',
|
||||
'EmotionResult',
|
||||
'FairFace',
|
||||
@@ -107,7 +114,7 @@ __all__ = [
|
||||
'BYTETracker',
|
||||
# Privacy
|
||||
'BlurFace',
|
||||
# Indexing (optional)
|
||||
# Stores (optional)
|
||||
'FAISS',
|
||||
# Utilities
|
||||
'Logger',
|
||||
|
||||
@@ -4,10 +4,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from uniface.attribute.age_gender import AgeGender
|
||||
from uniface.attribute.fairface import FairFace
|
||||
from uniface.attribute.base import Attribute
|
||||
from uniface.detection.base import BaseDetector
|
||||
from uniface.log import Logger
|
||||
from uniface.recognition.base import BaseRecognizer
|
||||
@@ -15,53 +16,73 @@ from uniface.types import Face
|
||||
|
||||
__all__ = ['FaceAnalyzer']
|
||||
|
||||
_UNSET: Any = object()
|
||||
|
||||
|
||||
class FaceAnalyzer:
|
||||
"""Unified face analyzer combining detection, recognition, and attributes.
|
||||
|
||||
This class provides a high-level interface for face analysis by combining
|
||||
multiple components: face detection, recognition (embedding extraction),
|
||||
and attribute prediction (age, gender, race).
|
||||
and an extensible list of attribute predictors (age, gender, race,
|
||||
emotion, etc.).
|
||||
|
||||
Any :class:`~uniface.attribute.base.Attribute` subclass can be passed
|
||||
via the ``attributes`` list. Each predictor's ``predict(image, face)``
|
||||
is called once per detected face, enriching the :class:`Face` in-place.
|
||||
|
||||
When called with no arguments, uses SCRFD (500M) for detection and
|
||||
ArcFace (MobileNet) for recognition — the smallest and fastest variants.
|
||||
|
||||
Args:
|
||||
detector: Face detector instance for detecting faces in images.
|
||||
recognizer: Optional face recognizer for extracting embeddings.
|
||||
age_gender: Optional age/gender predictor.
|
||||
fairface: Optional FairFace predictor for demographics.
|
||||
detector: Face detector instance. Defaults to ``SCRFD(SCRFD_500M_KPS)``.
|
||||
recognizer: Face recognizer for extracting embeddings.
|
||||
Defaults to ``ArcFace(MNET)``. Pass ``None`` to disable recognition.
|
||||
attributes: Optional list of ``Attribute`` predictors to run on
|
||||
each detected face (e.g. ``[AgeGender()]``).
|
||||
|
||||
Example:
|
||||
>>> from uniface import RetinaFace, ArcFace, FaceAnalyzer
|
||||
>>> detector = RetinaFace()
|
||||
>>> recognizer = ArcFace()
|
||||
>>> analyzer = FaceAnalyzer(detector, recognizer=recognizer)
|
||||
Examples:
|
||||
>>> from uniface import FaceAnalyzer
|
||||
>>> analyzer = FaceAnalyzer()
|
||||
>>> faces = analyzer.analyze(image)
|
||||
|
||||
>>> from uniface import FaceAnalyzer, AgeGender
|
||||
>>> analyzer = FaceAnalyzer(attributes=[AgeGender()])
|
||||
>>> faces = analyzer.analyze(image)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
detector: BaseDetector,
|
||||
recognizer: BaseRecognizer | None = None,
|
||||
age_gender: AgeGender | None = None,
|
||||
fairface: FairFace | None = None,
|
||||
detector: BaseDetector | None = None,
|
||||
recognizer: BaseRecognizer | None = _UNSET,
|
||||
attributes: list[Attribute] | None = None,
|
||||
) -> None:
|
||||
if detector is None:
|
||||
from uniface.constants import SCRFDWeights
|
||||
from uniface.detection import SCRFD
|
||||
|
||||
detector = SCRFD(model_name=SCRFDWeights.SCRFD_500M_KPS)
|
||||
|
||||
if recognizer is _UNSET:
|
||||
from uniface.recognition import ArcFace
|
||||
|
||||
recognizer = ArcFace()
|
||||
|
||||
self.detector = detector
|
||||
self.recognizer = recognizer
|
||||
self.age_gender = age_gender
|
||||
self.fairface = fairface
|
||||
self.attributes: list[Attribute] = attributes or []
|
||||
|
||||
Logger.info(f'Initialized FaceAnalyzer with detector={detector.__class__.__name__}')
|
||||
if recognizer:
|
||||
Logger.info(f' - Recognition enabled: {recognizer.__class__.__name__}')
|
||||
if age_gender:
|
||||
Logger.info(f' - Age/Gender enabled: {age_gender.__class__.__name__}')
|
||||
if fairface:
|
||||
Logger.info(f' - FairFace enabled: {fairface.__class__.__name__}')
|
||||
Logger.info(f'Recognition enabled: {recognizer.__class__.__name__}')
|
||||
for attr in self.attributes:
|
||||
Logger.info(f'Attribute enabled: {attr.__class__.__name__}')
|
||||
|
||||
def analyze(self, image: np.ndarray) -> list[Face]:
|
||||
"""Analyze faces in an image.
|
||||
|
||||
Performs face detection and optionally extracts embeddings and
|
||||
predicts attributes for each detected face.
|
||||
Performs face detection, optionally extracts embeddings, and runs
|
||||
every registered attribute predictor on each detected face.
|
||||
|
||||
Args:
|
||||
image: Input image as numpy array with shape (H, W, C) in BGR format.
|
||||
@@ -76,28 +97,17 @@ class FaceAnalyzer:
|
||||
if self.recognizer is not None:
|
||||
try:
|
||||
face.embedding = self.recognizer.get_normalized_embedding(image, face.landmarks)
|
||||
Logger.debug(f' Face {idx + 1}: Extracted embedding with shape {face.embedding.shape}')
|
||||
Logger.debug(f'Face {idx + 1}: Extracted embedding with shape {face.embedding.shape}')
|
||||
except Exception as e:
|
||||
Logger.warning(f' Face {idx + 1}: Failed to extract embedding: {e}')
|
||||
Logger.warning(f'Face {idx + 1}: Failed to extract embedding: {e}')
|
||||
|
||||
if self.age_gender is not None:
|
||||
for attr in self.attributes:
|
||||
attr_name = attr.__class__.__name__
|
||||
try:
|
||||
result = self.age_gender.predict(image, face.bbox)
|
||||
face.gender = result.gender
|
||||
face.age = result.age
|
||||
Logger.debug(f' Face {idx + 1}: Age={face.age}, Gender={face.sex}')
|
||||
attr.predict(image, face)
|
||||
Logger.debug(f'Face {idx + 1}: {attr_name} prediction succeeded')
|
||||
except Exception as e:
|
||||
Logger.warning(f' Face {idx + 1}: Failed to predict age/gender: {e}')
|
||||
|
||||
if self.fairface is not None:
|
||||
try:
|
||||
result = self.fairface.predict(image, face.bbox)
|
||||
face.gender = result.gender
|
||||
face.age_group = result.age_group
|
||||
face.race = result.race
|
||||
Logger.debug(f' Face {idx + 1}: AgeGroup={face.age_group}, Gender={face.sex}, Race={face.race}')
|
||||
except Exception as e:
|
||||
Logger.warning(f' Face {idx + 1}: Failed to predict FairFace attributes: {e}')
|
||||
Logger.warning(f'Face {idx + 1}: {attr_name} prediction failed: {e}')
|
||||
|
||||
Logger.info(f'Analysis complete: {len(faces)} face(s) processed')
|
||||
return faces
|
||||
@@ -106,8 +116,6 @@ class FaceAnalyzer:
|
||||
parts = [f'FaceAnalyzer(detector={self.detector.__class__.__name__}']
|
||||
if self.recognizer:
|
||||
parts.append(f'recognizer={self.recognizer.__class__.__name__}')
|
||||
if self.age_gender:
|
||||
parts.append(f'age_gender={self.age_gender.__class__.__name__}')
|
||||
if self.fairface:
|
||||
parts.append(f'fairface={self.fairface.__class__.__name__}')
|
||||
for attr in self.attributes:
|
||||
parts.append(f'{attr.__class__.__name__}')
|
||||
return ', '.join(parts) + ')'
|
||||
|
||||
@@ -12,7 +12,7 @@ from uniface.attribute.age_gender import AgeGender
|
||||
from uniface.attribute.base import Attribute
|
||||
from uniface.attribute.fairface import FairFace
|
||||
from uniface.constants import AgeGenderWeights, DDAMFNWeights, FairFaceWeights
|
||||
from uniface.types import AttributeResult, EmotionResult
|
||||
from uniface.types import AttributeResult, EmotionResult, Face
|
||||
|
||||
try:
|
||||
from uniface.attribute.emotion import Emotion
|
||||
@@ -30,7 +30,7 @@ except ImportError:
|
||||
def _initialize_model(self) -> None: ...
|
||||
def preprocess(self, image: np.ndarray, *args: Any) -> Any: ...
|
||||
def postprocess(self, prediction: Any) -> Any: ...
|
||||
def predict(self, image: np.ndarray, *args: Any) -> Any: ...
|
||||
def predict(self, image: np.ndarray, face: Face) -> Any: ...
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
||||
@@ -12,7 +12,7 @@ from uniface.face_utils import bbox_center_alignment
|
||||
from uniface.log import Logger
|
||||
from uniface.model_store import verify_model_weights
|
||||
from uniface.onnx_utils import create_onnx_session
|
||||
from uniface.types import AttributeResult
|
||||
from uniface.types import AttributeResult, Face
|
||||
|
||||
__all__ = ['AgeGender']
|
||||
|
||||
@@ -133,17 +133,20 @@ class AgeGender(Attribute):
|
||||
age = int(np.round(prediction[2] * 100))
|
||||
return AttributeResult(gender=gender, age=age)
|
||||
|
||||
def predict(self, image: np.ndarray, bbox: list | np.ndarray) -> AttributeResult:
|
||||
"""
|
||||
Predicts age and gender for a single face specified by a bounding box.
|
||||
def predict(self, image: np.ndarray, face: Face) -> AttributeResult:
|
||||
"""Predict age and gender and enrich the Face in-place.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The full input image in BGR format.
|
||||
bbox (Union[List, np.ndarray]): The face bounding box coordinates [x1, y1, x2, y2].
|
||||
image: The full input image in BGR format.
|
||||
face: Detected face; ``face.bbox`` is used for alignment.
|
||||
|
||||
Returns:
|
||||
AttributeResult: Result containing gender (0=Female, 1=Male) and age (in years).
|
||||
``AttributeResult`` with gender (0=Female, 1=Male) and age (years).
|
||||
"""
|
||||
face_blob = self.preprocess(image, bbox)
|
||||
face_blob = self.preprocess(image, face.bbox)
|
||||
prediction = self.session.run(self.output_names, {self.input_name: face_blob})[0][0]
|
||||
return self.postprocess(prediction)
|
||||
result = self.postprocess(prediction)
|
||||
|
||||
face.gender = result.gender
|
||||
face.age = result.age
|
||||
return result
|
||||
|
||||
@@ -2,95 +2,78 @@
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from uniface.types import AttributeResult, EmotionResult
|
||||
from uniface.types import AttributeResult, EmotionResult, Face
|
||||
|
||||
__all__ = ['Attribute', 'AttributeResult', 'EmotionResult']
|
||||
|
||||
|
||||
class Attribute(ABC):
|
||||
"""
|
||||
Abstract base class for face attribute models.
|
||||
"""Abstract base class for face attribute models.
|
||||
|
||||
This class defines the common interface that all attribute models
|
||||
(e.g., age-gender, emotion) must implement. It ensures a consistent API
|
||||
across different attribute prediction modules in the library, making them
|
||||
interchangeable and easy to use.
|
||||
All attribute models (age-gender, emotion, FairFace, etc.) implement this
|
||||
interface so they can be used interchangeably inside ``FaceAnalyzer``.
|
||||
|
||||
The ``predict`` method accepts an image and a :class:`Face` object. Each
|
||||
subclass extracts what it needs (bbox, landmarks) from the Face, runs
|
||||
inference, writes the results back to the Face **and** returns a typed
|
||||
result dataclass.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def _initialize_model(self) -> None:
|
||||
"""
|
||||
Initializes the underlying model for inference.
|
||||
|
||||
This method should handle loading model weights, creating the
|
||||
inference session (e.g., ONNX Runtime, PyTorch), and any necessary
|
||||
warm-up procedures to prepare the model for prediction.
|
||||
"""
|
||||
"""Load model weights and create the inference session."""
|
||||
raise NotImplementedError('Subclasses must implement the _initialize_model method.')
|
||||
|
||||
@abstractmethod
|
||||
def preprocess(self, image: np.ndarray, *args: Any) -> Any:
|
||||
"""
|
||||
Preprocesses the input data for the model.
|
||||
|
||||
This method should take a raw image and any other necessary data
|
||||
(like bounding boxes or landmarks) and convert it into the format
|
||||
expected by the model's inference engine (e.g., a blob or tensor).
|
||||
"""Preprocess the input data for the model.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The input image containing the face, typically
|
||||
in BGR format.
|
||||
*args: Additional arguments required for preprocessing, such as
|
||||
bounding boxes or facial landmarks.
|
||||
image: The input image in BGR format.
|
||||
*args: Subclass-specific data (bbox, landmarks, etc.).
|
||||
|
||||
Returns:
|
||||
The preprocessed data ready for model inference.
|
||||
Preprocessed data ready for model inference.
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the preprocess method.')
|
||||
|
||||
@abstractmethod
|
||||
def postprocess(self, prediction: Any) -> Any:
|
||||
"""
|
||||
Postprocesses the raw model output into a human-readable format.
|
||||
|
||||
This method takes the raw output from the model's inference and
|
||||
converts it into a meaningful result, such as an age value, a gender
|
||||
label, or an emotion category.
|
||||
"""Convert raw model output into a typed result dataclass.
|
||||
|
||||
Args:
|
||||
prediction (Any): The raw output from the model's inference.
|
||||
prediction: Raw output from the model.
|
||||
|
||||
Returns:
|
||||
The final, processed attributes.
|
||||
An ``AttributeResult`` or ``EmotionResult``.
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the postprocess method.')
|
||||
|
||||
@abstractmethod
|
||||
def predict(self, image: np.ndarray, *args: Any) -> Any:
|
||||
"""
|
||||
Performs end-to-end attribute prediction on a given image.
|
||||
def predict(self, image: np.ndarray, face: Face) -> AttributeResult | EmotionResult:
|
||||
"""Run end-to-end prediction and enrich the Face in-place.
|
||||
|
||||
This method orchestrates the full pipeline: it calls the preprocess,
|
||||
inference, and postprocess steps to return the final, user-friendly
|
||||
attribute prediction.
|
||||
Each subclass extracts what it needs from *face* (e.g. ``face.bbox``
|
||||
or ``face.landmarks``), runs the full preprocess-infer-postprocess
|
||||
pipeline, writes relevant fields back to *face*, and returns the
|
||||
result dataclass.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The input image containing the face.
|
||||
*args: Additional data required for prediction, such as a bounding
|
||||
box or landmarks.
|
||||
image: The full input image in BGR format.
|
||||
face: Detected face whose attribute fields will be populated.
|
||||
|
||||
Returns:
|
||||
The final predicted attributes.
|
||||
The prediction result (``AttributeResult`` or ``EmotionResult``).
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the predict method.')
|
||||
|
||||
def __call__(self, *args, **kwargs) -> Any:
|
||||
"""
|
||||
Provides a convenient, callable shortcut for the `predict` method.
|
||||
"""
|
||||
return self.predict(*args, **kwargs)
|
||||
def __call__(self, image: np.ndarray, face: Face) -> AttributeResult | EmotionResult:
|
||||
"""Callable shortcut for :meth:`predict`."""
|
||||
return self.predict(image, face)
|
||||
|
||||
@@ -12,7 +12,7 @@ from uniface.constants import DDAMFNWeights
|
||||
from uniface.face_utils import face_alignment
|
||||
from uniface.log import Logger
|
||||
from uniface.model_store import verify_model_weights
|
||||
from uniface.types import EmotionResult
|
||||
from uniface.types import EmotionResult, Face
|
||||
|
||||
__all__ = ['Emotion']
|
||||
|
||||
@@ -116,14 +116,23 @@ class Emotion(Attribute):
|
||||
confidence = float(probabilities[pred_index])
|
||||
return EmotionResult(emotion=emotion_label, confidence=confidence)
|
||||
|
||||
def predict(self, image: np.ndarray, landmark: list | np.ndarray) -> EmotionResult:
|
||||
def predict(self, image: np.ndarray, face: Face) -> EmotionResult:
|
||||
"""Predict emotion and enrich the Face in-place.
|
||||
|
||||
Args:
|
||||
image: The full input image in BGR format.
|
||||
face: Detected face; ``face.landmarks`` is used for alignment.
|
||||
|
||||
Returns:
|
||||
``EmotionResult`` with emotion label and confidence score.
|
||||
"""
|
||||
Predicts the emotion from a single face specified by its landmarks.
|
||||
"""
|
||||
input_tensor = self.preprocess(image, landmark)
|
||||
input_tensor = self.preprocess(image, face.landmarks)
|
||||
with torch.no_grad():
|
||||
output = self.model(input_tensor)
|
||||
if isinstance(output, tuple):
|
||||
output = output[0]
|
||||
|
||||
return self.postprocess(output)
|
||||
result = self.postprocess(output)
|
||||
face.emotion = result.emotion
|
||||
face.emotion_confidence = result.confidence
|
||||
return result
|
||||
|
||||
@@ -11,7 +11,7 @@ from uniface.constants import FairFaceWeights
|
||||
from uniface.log import Logger
|
||||
from uniface.model_store import verify_model_weights
|
||||
from uniface.onnx_utils import create_onnx_session
|
||||
from uniface.types import AttributeResult
|
||||
from uniface.types import AttributeResult, Face
|
||||
|
||||
__all__ = ['AGE_LABELS', 'RACE_LABELS', 'FairFace']
|
||||
|
||||
@@ -168,29 +168,24 @@ class FairFace(Attribute):
|
||||
race=RACE_LABELS[race_idx],
|
||||
)
|
||||
|
||||
def predict(self, image: np.ndarray, bbox: list | np.ndarray | None = None) -> AttributeResult:
|
||||
"""
|
||||
Predicts race, gender, and age for a face.
|
||||
def predict(self, image: np.ndarray, face: Face) -> AttributeResult:
|
||||
"""Predict race, gender, and age and enrich the Face in-place.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The input image in BGR format.
|
||||
bbox (Optional[Union[List, np.ndarray]]): Face bounding box [x1, y1, x2, y2].
|
||||
If None, uses the entire image.
|
||||
image: The full input image in BGR format.
|
||||
face: Detected face; ``face.bbox`` is used for cropping.
|
||||
|
||||
Returns:
|
||||
AttributeResult: Result containing:
|
||||
- gender: 0=Female, 1=Male
|
||||
- age_group: Age range string like "20-29"
|
||||
- race: Race/ethnicity label
|
||||
``AttributeResult`` with gender, age_group, and race.
|
||||
"""
|
||||
# Preprocess
|
||||
input_blob = self.preprocess(image, bbox)
|
||||
|
||||
# Inference
|
||||
input_blob = self.preprocess(image, face.bbox)
|
||||
outputs = self.session.run(self.output_names, {self.input_name: input_blob})
|
||||
result = self.postprocess(outputs)
|
||||
|
||||
# Postprocess
|
||||
return self.postprocess(outputs)
|
||||
face.gender = result.gender
|
||||
face.age_group = result.age_group
|
||||
face.race = result.race
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _softmax(x: np.ndarray) -> np.ndarray:
|
||||
|
||||
@@ -16,6 +16,7 @@ __all__ = [
|
||||
'distance2bbox',
|
||||
'distance2kps',
|
||||
'generate_anchors',
|
||||
'letterbox_resize',
|
||||
'non_max_suppression',
|
||||
'resize_image',
|
||||
'xyxy_to_cxcywh',
|
||||
@@ -277,3 +278,70 @@ def distance2kps(
|
||||
preds.append(px)
|
||||
preds.append(py)
|
||||
return np.stack(preds, axis=-1)
|
||||
|
||||
|
||||
def letterbox_resize(
|
||||
image: np.ndarray,
|
||||
target_size: int,
|
||||
fill_value: int = 114,
|
||||
) -> tuple[np.ndarray, float, tuple[int, int]]:
|
||||
"""Letterbox resize with center padding for YOLO-style detectors.
|
||||
|
||||
Maintains aspect ratio by scaling the image to fit within target_size,
|
||||
then center-pads with a constant fill value. Converts BGR to RGB,
|
||||
normalizes to [0, 1], and transposes to NCHW format.
|
||||
|
||||
This preprocessing strategy is standard for YOLO models and ensures
|
||||
no distortion while maintaining a square input size.
|
||||
|
||||
Args:
|
||||
image: Input image in BGR format with shape (H, W, C).
|
||||
target_size: Target square size (e.g., 640 for 640x640 input).
|
||||
fill_value: Padding fill value (default: 114 for gray background).
|
||||
|
||||
Returns:
|
||||
Tuple of (preprocessed_tensor, scale_ratio, padding):
|
||||
- preprocessed_tensor: Shape (1, 3, target_size, target_size),
|
||||
RGB, normalized [0, 1], NCHW format, float32, contiguous.
|
||||
- scale_ratio: Resize scale factor for coordinate transformation.
|
||||
- padding: Padding offsets as (pad_w, pad_h) for coordinate transformation.
|
||||
|
||||
Example:
|
||||
>>> image = cv2.imread('face.jpg') # (480, 640, 3)
|
||||
>>> tensor, scale, (pad_w, pad_h) = letterbox_resize(image, 640)
|
||||
>>> tensor.shape
|
||||
(1, 3, 640, 640)
|
||||
>>> # To transform coordinates back to original:
|
||||
>>> x_orig = (x_detected - pad_w) / scale
|
||||
>>> y_orig = (y_detected - pad_h) / scale
|
||||
"""
|
||||
# Get original image shape
|
||||
img_h, img_w = image.shape[:2]
|
||||
|
||||
# Calculate scale ratio to fit within target_size
|
||||
scale = min(target_size / img_h, target_size / img_w)
|
||||
new_h, new_w = int(img_h * scale), int(img_w * scale)
|
||||
|
||||
# Resize image maintaining aspect ratio
|
||||
img_resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Create padded canvas with fill_value
|
||||
img_padded = np.full((target_size, target_size, 3), fill_value, dtype=np.uint8)
|
||||
|
||||
# Calculate padding to center the image
|
||||
pad_h = (target_size - new_h) // 2
|
||||
pad_w = (target_size - new_w) // 2
|
||||
|
||||
# Place resized image in center of canvas
|
||||
img_padded[pad_h : pad_h + new_h, pad_w : pad_w + new_w] = img_resized
|
||||
|
||||
# Convert BGR to RGB and normalize to [0, 1]
|
||||
img_rgb = cv2.cvtColor(img_padded, cv2.COLOR_BGR2RGB)
|
||||
img_normalized = img_rgb.astype(np.float32) / 255.0
|
||||
|
||||
# Transpose to CHW format and add batch dimension (NCHW)
|
||||
img_transposed = np.transpose(img_normalized, (2, 0, 1))
|
||||
img_batch = np.expand_dims(img_transposed, axis=0)
|
||||
img_batch = np.ascontiguousarray(img_batch)
|
||||
|
||||
return img_batch, scale, (pad_w, pad_h)
|
||||
|
||||
@@ -156,6 +156,20 @@ class GazeWeights(str, Enum):
|
||||
MOBILEONE_S0 = "gaze_mobileone_s0"
|
||||
|
||||
|
||||
class HeadPoseWeights(str, Enum):
|
||||
"""
|
||||
Head pose estimation models using 6D rotation representation.
|
||||
Trained on 300W-LP dataset, evaluated on AFLW2000.
|
||||
https://github.com/yakhyo/head-pose-estimation
|
||||
"""
|
||||
RESNET18 = "headpose_resnet18"
|
||||
RESNET34 = "headpose_resnet34"
|
||||
RESNET50 = "headpose_resnet50"
|
||||
MOBILENET_V2 = "headpose_mobilenetv2"
|
||||
MOBILENET_V3_SMALL = "headpose_mobilenetv3_small"
|
||||
MOBILENET_V3_LARGE = "headpose_mobilenetv3_large"
|
||||
|
||||
|
||||
class ParsingWeights(str, Enum):
|
||||
"""
|
||||
Face Parsing: Semantic Segmentation of Facial Components.
|
||||
@@ -348,6 +362,32 @@ MODEL_REGISTRY: dict[Enum, ModelInfo] = {
|
||||
sha256='8b4fdc4e3da44733c9a82e7776b411e4a39f94e8e285aee0fc85a548a55f7d9f'
|
||||
),
|
||||
|
||||
# Head Pose
|
||||
HeadPoseWeights.RESNET18: ModelInfo(
|
||||
url='https://github.com/yakhyo/head-pose-estimation/releases/download/weights/resnet18.onnx',
|
||||
sha256='61c34e877989412980d1ea80c52391250b074abc00d19a6100de5c8e999212ee'
|
||||
),
|
||||
HeadPoseWeights.RESNET34: ModelInfo(
|
||||
url='https://github.com/yakhyo/head-pose-estimation/releases/download/weights/resnet34.onnx',
|
||||
sha256='8da9f2ce4810298ebea68bd85fba1b6bd11716060c10534596f46be52cc908c9'
|
||||
),
|
||||
HeadPoseWeights.RESNET50: ModelInfo(
|
||||
url='https://github.com/yakhyo/head-pose-estimation/releases/download/weights/resnet50.onnx',
|
||||
sha256='50c74d57b7663361b8ede83b0e4122546171119ef502ec55b790dbd7fc360260'
|
||||
),
|
||||
HeadPoseWeights.MOBILENET_V2: ModelInfo(
|
||||
url='https://github.com/yakhyo/head-pose-estimation/releases/download/weights/mobilenetv2.onnx',
|
||||
sha256='1e902872868e483bd0e4f8f4a8ff2a4d61c2ccbca9dadf748e5479b5cc86a9e9'
|
||||
),
|
||||
HeadPoseWeights.MOBILENET_V3_SMALL: ModelInfo(
|
||||
url='https://github.com/yakhyo/head-pose-estimation/releases/download/weights/mobilenetv3_small.onnx',
|
||||
sha256='e8ae4d932b3d13221638fc72e171603e020c6da28b770753f76146867f40e190'
|
||||
),
|
||||
HeadPoseWeights.MOBILENET_V3_LARGE: ModelInfo(
|
||||
url='https://github.com/yakhyo/head-pose-estimation/releases/download/weights/mobilenetv3_large.onnx',
|
||||
sha256='3a68815fa00aba41ddc4e014bf631b637caba8619df71160383f1fee8c15a3c9'
|
||||
),
|
||||
|
||||
# Parsing
|
||||
ParsingWeights.RESNET18: ModelInfo(
|
||||
url='https://github.com/yakhyo/face-parsing/releases/download/weights/resnet18.onnx',
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -119,3 +119,77 @@ class BaseDetector(ABC):
|
||||
List of detected Face objects.
|
||||
"""
|
||||
return self.detect(image, **kwargs)
|
||||
|
||||
def _select_top_detections(
|
||||
self,
|
||||
detections: np.ndarray,
|
||||
landmarks: np.ndarray,
|
||||
max_num: int,
|
||||
original_shape: tuple[int, int],
|
||||
metric: Literal['default', 'max'] = 'max',
|
||||
center_weight: float = 2.0,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Filter detections to keep only top max_num faces.
|
||||
|
||||
Ranks faces by area and/or distance from image center, then selects
|
||||
the top max_num detections.
|
||||
|
||||
Args:
|
||||
detections: Array of shape (N, 5) as [x1, y1, x2, y2, confidence].
|
||||
landmarks: Array of shape (N, 5, 2) for 5-point landmarks.
|
||||
max_num: Maximum number of faces to keep. If 0 or >= N, returns all.
|
||||
original_shape: Original image shape as (height, width).
|
||||
metric: Ranking metric:
|
||||
- 'max': Rank by bounding box area only.
|
||||
- 'default': Rank by area minus center distance penalty.
|
||||
center_weight: Weight for center distance penalty (only used with 'default' metric).
|
||||
|
||||
Returns:
|
||||
Filtered (detections, landmarks) tuple with at most max_num faces.
|
||||
"""
|
||||
if max_num <= 0 or detections.shape[0] <= max_num:
|
||||
return detections, landmarks
|
||||
|
||||
# Calculate bounding box areas
|
||||
area = (detections[:, 2] - detections[:, 0]) * (detections[:, 3] - detections[:, 1])
|
||||
|
||||
# Calculate offsets from image center
|
||||
center_y, center_x = original_shape[0] // 2, original_shape[1] // 2
|
||||
offsets = np.vstack(
|
||||
[
|
||||
(detections[:, 0] + detections[:, 2]) / 2 - center_x,
|
||||
(detections[:, 1] + detections[:, 3]) / 2 - center_y,
|
||||
]
|
||||
)
|
||||
offset_dist_squared = np.sum(np.power(offsets, 2.0), axis=0)
|
||||
|
||||
# Calculate ranking scores based on metric
|
||||
if metric == 'max':
|
||||
scores = area
|
||||
else:
|
||||
scores = area - offset_dist_squared * center_weight
|
||||
|
||||
# Select top max_num by score
|
||||
top_indices = np.argsort(scores)[::-1][:max_num]
|
||||
return detections[top_indices], landmarks[top_indices]
|
||||
|
||||
@staticmethod
|
||||
def _detections_to_faces(detections: np.ndarray, landmarks: np.ndarray) -> list[Face]:
|
||||
"""Convert detection arrays to Face objects.
|
||||
|
||||
Args:
|
||||
detections: Array of shape (N, 5) as [x1, y1, x2, y2, confidence].
|
||||
landmarks: Array of shape (N, 5, 2) for 5-point landmarks.
|
||||
|
||||
Returns:
|
||||
List of Face objects.
|
||||
"""
|
||||
faces = []
|
||||
for i in range(detections.shape[0]):
|
||||
face = Face(
|
||||
bbox=detections[i, :4],
|
||||
confidence=float(detections[i, 4]),
|
||||
landmarks=landmarks[i],
|
||||
)
|
||||
faces.append(face)
|
||||
return faces
|
||||
|
||||
@@ -208,42 +208,12 @@ class RetinaFace(BaseDetector):
|
||||
# Postprocessing
|
||||
detections, landmarks = self.postprocess(outputs, resize_factor, shape=(width, height))
|
||||
|
||||
if max_num > 0 and detections.shape[0] > max_num:
|
||||
# Calculate area of detections
|
||||
areas = (detections[:, 2] - detections[:, 0]) * (detections[:, 3] - detections[:, 1])
|
||||
# Filter to top max_num faces if requested
|
||||
detections, landmarks = self._select_top_detections(
|
||||
detections, landmarks, max_num, (original_height, original_width), metric, center_weight
|
||||
)
|
||||
|
||||
# Calculate offsets from image center
|
||||
center = (original_height // 2, original_width // 2)
|
||||
offsets = np.vstack(
|
||||
[
|
||||
(detections[:, 0] + detections[:, 2]) / 2 - center[1],
|
||||
(detections[:, 1] + detections[:, 3]) / 2 - center[0],
|
||||
]
|
||||
)
|
||||
offset_dist_squared = np.sum(np.power(offsets, 2.0), axis=0)
|
||||
|
||||
# Calculate scores based on the chosen metric
|
||||
if metric == 'max':
|
||||
scores = areas
|
||||
else:
|
||||
scores = areas - offset_dist_squared * center_weight
|
||||
|
||||
# Sort by scores and select top `max_num`
|
||||
sorted_indices = np.argsort(scores)[::-1][:max_num]
|
||||
|
||||
detections = detections[sorted_indices]
|
||||
landmarks = landmarks[sorted_indices]
|
||||
|
||||
faces = []
|
||||
for i in range(detections.shape[0]):
|
||||
face = Face(
|
||||
bbox=detections[i, :4],
|
||||
confidence=float(detections[i, 4]),
|
||||
landmarks=landmarks[i],
|
||||
)
|
||||
faces.append(face)
|
||||
|
||||
return faces
|
||||
return self._detections_to_faces(detections, landmarks)
|
||||
|
||||
def postprocess(
|
||||
self,
|
||||
|
||||
@@ -272,38 +272,9 @@ class SCRFD(BaseDetector):
|
||||
landmarks = landmarks[order, :, :]
|
||||
landmarks = landmarks[keep, :, :].astype(np.float32)
|
||||
|
||||
if 0 < max_num < detections.shape[0]:
|
||||
# Calculate area of detections
|
||||
area = (detections[:, 2] - detections[:, 0]) * (detections[:, 3] - detections[:, 1])
|
||||
# Filter to top max_num faces if requested
|
||||
detections, landmarks = self._select_top_detections(
|
||||
detections, landmarks, max_num, (original_height, original_width), metric, center_weight
|
||||
)
|
||||
|
||||
# Calculate offsets from image center
|
||||
center = (original_height // 2, original_width // 2)
|
||||
offsets = np.vstack(
|
||||
[
|
||||
(detections[:, 0] + detections[:, 2]) / 2 - center[1],
|
||||
(detections[:, 1] + detections[:, 3]) / 2 - center[0],
|
||||
]
|
||||
)
|
||||
|
||||
# Calculate scores based on the chosen metric
|
||||
offset_dist_squared = np.sum(np.power(offsets, 2.0), axis=0)
|
||||
if metric == 'max':
|
||||
values = area
|
||||
else:
|
||||
values = area - offset_dist_squared * center_weight
|
||||
|
||||
# Sort by scores and select top `max_num`
|
||||
sorted_indices = np.argsort(values)[::-1][:max_num]
|
||||
detections = detections[sorted_indices]
|
||||
landmarks = landmarks[sorted_indices]
|
||||
|
||||
faces = []
|
||||
for i in range(detections.shape[0]):
|
||||
face = Face(
|
||||
bbox=detections[i, :4],
|
||||
confidence=float(detections[i, 4]),
|
||||
landmarks=landmarks[i],
|
||||
)
|
||||
faces.append(face)
|
||||
|
||||
return faces
|
||||
return self._detections_to_faces(detections, landmarks)
|
||||
|
||||
@@ -4,10 +4,9 @@
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from uniface.common import non_max_suppression
|
||||
from uniface.common import letterbox_resize, non_max_suppression
|
||||
from uniface.constants import YOLOv5FaceWeights
|
||||
from uniface.log import Logger
|
||||
from uniface.model_store import verify_model_weights
|
||||
@@ -140,45 +139,15 @@ class YOLOv5Face(BaseDetector):
|
||||
raise RuntimeError(f"Failed to initialize model session for '{model_path}'") from e
|
||||
|
||||
def preprocess(self, image: np.ndarray) -> tuple[np.ndarray, float, tuple[int, int]]:
|
||||
"""
|
||||
Preprocess image for inference.
|
||||
"""Preprocess image using letterbox resize.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): Input image (BGR format)
|
||||
image: Input image in BGR format.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, float, Tuple[int, int]]: Preprocessed image, scale ratio, and padding
|
||||
Tuple of (preprocessed_tensor, scale_ratio, padding).
|
||||
"""
|
||||
# Get original image shape
|
||||
img_h, img_w = image.shape[:2]
|
||||
|
||||
# Calculate scale ratio
|
||||
scale = min(self.input_size / img_h, self.input_size / img_w)
|
||||
new_h, new_w = int(img_h * scale), int(img_w * scale)
|
||||
|
||||
# Resize image
|
||||
img_resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Create padded image
|
||||
img_padded = np.full((self.input_size, self.input_size, 3), 114, dtype=np.uint8)
|
||||
|
||||
# Calculate padding
|
||||
pad_h = (self.input_size - new_h) // 2
|
||||
pad_w = (self.input_size - new_w) // 2
|
||||
|
||||
# Place resized image in center
|
||||
img_padded[pad_h : pad_h + new_h, pad_w : pad_w + new_w] = img_resized
|
||||
|
||||
# Convert to RGB and normalize
|
||||
img_rgb = cv2.cvtColor(img_padded, cv2.COLOR_BGR2RGB)
|
||||
img_normalized = img_rgb.astype(np.float32) / 255.0
|
||||
|
||||
# Transpose to CHW format (HWC -> CHW) and add batch dimension
|
||||
img_transposed = np.transpose(img_normalized, (2, 0, 1))
|
||||
img_batch = np.expand_dims(img_transposed, axis=0)
|
||||
img_batch = np.ascontiguousarray(img_batch)
|
||||
|
||||
return img_batch, scale, (pad_w, pad_h)
|
||||
return letterbox_resize(image, self.input_size)
|
||||
|
||||
def inference(self, input_tensor: np.ndarray) -> list[np.ndarray]:
|
||||
"""Perform model inference on the preprocessed image tensor.
|
||||
@@ -337,38 +306,9 @@ class YOLOv5Face(BaseDetector):
|
||||
if len(detections) == 0:
|
||||
return []
|
||||
|
||||
if 0 < max_num < detections.shape[0]:
|
||||
# Calculate area of detections
|
||||
area = (detections[:, 2] - detections[:, 0]) * (detections[:, 3] - detections[:, 1])
|
||||
# Filter to top max_num faces if requested
|
||||
detections, landmarks = self._select_top_detections(
|
||||
detections, landmarks, max_num, (original_height, original_width), metric, center_weight
|
||||
)
|
||||
|
||||
# Calculate offsets from image center
|
||||
center = (original_height // 2, original_width // 2)
|
||||
offsets = np.vstack(
|
||||
[
|
||||
(detections[:, 0] + detections[:, 2]) / 2 - center[1],
|
||||
(detections[:, 1] + detections[:, 3]) / 2 - center[0],
|
||||
]
|
||||
)
|
||||
|
||||
# Calculate scores based on the chosen metric
|
||||
offset_dist_squared = np.sum(np.power(offsets, 2.0), axis=0)
|
||||
if metric == 'max':
|
||||
values = area
|
||||
else:
|
||||
values = area - offset_dist_squared * center_weight
|
||||
|
||||
# Sort by scores and select top `max_num`
|
||||
sorted_indices = np.argsort(values)[::-1][:max_num]
|
||||
detections = detections[sorted_indices]
|
||||
landmarks = landmarks[sorted_indices]
|
||||
|
||||
faces = []
|
||||
for i in range(detections.shape[0]):
|
||||
face = Face(
|
||||
bbox=detections[i, :4],
|
||||
confidence=float(detections[i, 4]),
|
||||
landmarks=landmarks[i],
|
||||
)
|
||||
faces.append(face)
|
||||
|
||||
return faces
|
||||
return self._detections_to_faces(detections, landmarks)
|
||||
|
||||
@@ -11,10 +11,9 @@ Reference: https://github.com/yakhyo/yolov8-face-onnx-inference
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from uniface.common import non_max_suppression
|
||||
from uniface.common import letterbox_resize, non_max_suppression
|
||||
from uniface.constants import YOLOv8FaceWeights
|
||||
from uniface.log import Logger
|
||||
from uniface.model_store import verify_model_weights
|
||||
@@ -151,45 +150,15 @@ class YOLOv8Face(BaseDetector):
|
||||
raise RuntimeError(f"Failed to initialize model session for '{model_path}'") from e
|
||||
|
||||
def preprocess(self, image: np.ndarray) -> tuple[np.ndarray, float, tuple[int, int]]:
|
||||
"""
|
||||
Preprocess image for inference (letterbox resize with center padding).
|
||||
"""Preprocess image using letterbox resize.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): Input image (BGR format)
|
||||
image: Input image in BGR format.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, float, Tuple[int, int]]: Preprocessed image, scale ratio, and padding (pad_w, pad_h)
|
||||
Tuple of (preprocessed_tensor, scale_ratio, padding).
|
||||
"""
|
||||
# Get original image shape
|
||||
img_h, img_w = image.shape[:2]
|
||||
|
||||
# Calculate scale ratio
|
||||
scale = min(self.input_size / img_h, self.input_size / img_w)
|
||||
new_h, new_w = int(img_h * scale), int(img_w * scale)
|
||||
|
||||
# Resize image
|
||||
img_resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Create padded image with gray background (114, 114, 114)
|
||||
img_padded = np.full((self.input_size, self.input_size, 3), 114, dtype=np.uint8)
|
||||
|
||||
# Calculate padding (center the image)
|
||||
pad_h = (self.input_size - new_h) // 2
|
||||
pad_w = (self.input_size - new_w) // 2
|
||||
|
||||
# Place resized image in center
|
||||
img_padded[pad_h : pad_h + new_h, pad_w : pad_w + new_w] = img_resized
|
||||
|
||||
# Convert BGR to RGB and normalize
|
||||
img_rgb = cv2.cvtColor(img_padded, cv2.COLOR_BGR2RGB)
|
||||
img_normalized = img_rgb.astype(np.float32) / 255.0
|
||||
|
||||
# Transpose to CHW format (HWC -> CHW) and add batch dimension
|
||||
img_transposed = np.transpose(img_normalized, (2, 0, 1))
|
||||
img_batch = np.expand_dims(img_transposed, axis=0)
|
||||
img_batch = np.ascontiguousarray(img_batch)
|
||||
|
||||
return img_batch, scale, (pad_w, pad_h)
|
||||
return letterbox_resize(image, self.input_size)
|
||||
|
||||
def inference(self, input_tensor: np.ndarray) -> list[np.ndarray]:
|
||||
"""Perform model inference on the preprocessed image tensor.
|
||||
@@ -387,38 +356,9 @@ class YOLOv8Face(BaseDetector):
|
||||
if len(detections) == 0:
|
||||
return []
|
||||
|
||||
if 0 < max_num < detections.shape[0]:
|
||||
# Calculate area of detections
|
||||
area = (detections[:, 2] - detections[:, 0]) * (detections[:, 3] - detections[:, 1])
|
||||
# Filter to top max_num faces if requested
|
||||
detections, landmarks = self._select_top_detections(
|
||||
detections, landmarks, max_num, (original_height, original_width), metric, center_weight
|
||||
)
|
||||
|
||||
# Calculate offsets from image center
|
||||
center = (original_height // 2, original_width // 2)
|
||||
offsets = np.vstack(
|
||||
[
|
||||
(detections[:, 0] + detections[:, 2]) / 2 - center[1],
|
||||
(detections[:, 1] + detections[:, 3]) / 2 - center[0],
|
||||
]
|
||||
)
|
||||
|
||||
# Calculate scores based on the chosen metric
|
||||
offset_dist_squared = np.sum(np.power(offsets, 2.0), axis=0)
|
||||
if metric == 'max':
|
||||
values = area
|
||||
else:
|
||||
values = area - offset_dist_squared * center_weight
|
||||
|
||||
# Sort by scores and select top `max_num`
|
||||
sorted_indices = np.argsort(values)[::-1][:max_num]
|
||||
detections = detections[sorted_indices]
|
||||
landmarks = landmarks[sorted_indices]
|
||||
|
||||
faces = []
|
||||
for i in range(detections.shape[0]):
|
||||
face = Face(
|
||||
bbox=detections[i, :4],
|
||||
confidence=float(detections[i, 4]),
|
||||
landmarks=landmarks[i],
|
||||
)
|
||||
faces.append(face)
|
||||
|
||||
return faces
|
||||
return self._detections_to_faces(detections, landmarks)
|
||||
|
||||
237
uniface/draw.py
237
uniface/draw.py
@@ -21,6 +21,9 @@ __all__ = [
|
||||
'draw_corner_bbox',
|
||||
'draw_detections',
|
||||
'draw_gaze',
|
||||
'draw_head_pose',
|
||||
'draw_head_pose_axis',
|
||||
'draw_head_pose_cube',
|
||||
'draw_text_label',
|
||||
'draw_tracks',
|
||||
'vis_parsing_maps',
|
||||
@@ -229,9 +232,10 @@ def draw_text_label(
|
||||
def draw_detections(
|
||||
*,
|
||||
image: np.ndarray,
|
||||
bboxes: list[np.ndarray] | list[list[float]],
|
||||
scores: np.ndarray | list[float],
|
||||
landmarks: list[np.ndarray] | list[list[list[float]]],
|
||||
faces: list[Face] | None = None,
|
||||
bboxes: list[np.ndarray] | list[list[float]] | None = None,
|
||||
scores: np.ndarray | list[float] | None = None,
|
||||
landmarks: list[np.ndarray] | list[list[list[float]]] | None = None,
|
||||
vis_threshold: float = 0.6,
|
||||
draw_score: bool = False,
|
||||
corner_bbox: bool = True,
|
||||
@@ -240,17 +244,31 @@ def draw_detections(
|
||||
|
||||
Modifies the image in-place.
|
||||
|
||||
Accepts either a list of :class:`Face` objects (preferred) or separate
|
||||
lists of bboxes, scores, and landmarks for backward compatibility.
|
||||
|
||||
Args:
|
||||
image: Input image to draw on (modified in-place).
|
||||
faces: List of Face objects from detection. When provided,
|
||||
``bboxes``, ``scores``, and ``landmarks`` are ignored.
|
||||
bboxes: List of bounding boxes in xyxy format ``[x1, y1, x2, y2]``.
|
||||
scores: List of confidence scores.
|
||||
landmarks: List of landmark sets with shape ``(5, 2)``.
|
||||
vis_threshold: Confidence threshold for filtering. Defaults to 0.6.
|
||||
draw_score: Whether to draw confidence scores. Defaults to False.
|
||||
corner_bbox: Use corner-style bounding boxes. Defaults to True.
|
||||
"""
|
||||
|
||||
# Adaptive line thickness
|
||||
Examples:
|
||||
>>> draw_detections(image=image, faces=faces)
|
||||
>>> draw_detections(image=image, faces=faces, vis_threshold=0.7, draw_score=True)
|
||||
"""
|
||||
if faces is not None:
|
||||
bboxes = [f.bbox for f in faces]
|
||||
scores = [f.confidence for f in faces]
|
||||
landmarks = [f.landmarks for f in faces]
|
||||
elif bboxes is None or scores is None or landmarks is None:
|
||||
raise ValueError('Provide either faces or all of bboxes, scores, and landmarks')
|
||||
|
||||
line_thickness = max(round(sum(image.shape[:2]) / 2 * 0.003), 2)
|
||||
|
||||
for i, score in enumerate(scores):
|
||||
@@ -259,13 +277,11 @@ def draw_detections(
|
||||
|
||||
bbox = np.array(bboxes[i], dtype=np.int32)
|
||||
|
||||
# Draw bounding box
|
||||
if corner_bbox:
|
||||
draw_corner_bbox(image, bbox, color=(0, 255, 0), thickness=line_thickness, proportion=0.2)
|
||||
else:
|
||||
cv2.rectangle(image, tuple(bbox[:2]), tuple(bbox[2:]), (0, 255, 0), line_thickness)
|
||||
|
||||
# Draw confidence score label
|
||||
if draw_score:
|
||||
font_scale = max(0.4, min(0.7, (bbox[3] - bbox[1]) / 200))
|
||||
draw_text_label(
|
||||
@@ -278,7 +294,6 @@ def draw_detections(
|
||||
font_scale=font_scale,
|
||||
)
|
||||
|
||||
# Draw landmarks
|
||||
landmark_set = np.array(landmarks[i], dtype=np.int32)
|
||||
for j, point in enumerate(landmark_set):
|
||||
cv2.circle(image, tuple(point), line_thickness + 1, _LANDMARK_COLORS[j % len(_LANDMARK_COLORS)], -1)
|
||||
@@ -356,6 +371,212 @@ def draw_gaze(
|
||||
)
|
||||
|
||||
|
||||
def draw_head_pose_cube(
|
||||
image: np.ndarray,
|
||||
yaw: float,
|
||||
pitch: float,
|
||||
roll: float,
|
||||
bbox: list[int] | np.ndarray,
|
||||
size: int | None = None,
|
||||
) -> None:
|
||||
"""Draw a 3D wireframe cube representing head orientation on an image.
|
||||
|
||||
Projects a 3D cube onto the image plane based on yaw, pitch, and roll
|
||||
angles, centered on the face bounding box.
|
||||
|
||||
Modifies the image in-place.
|
||||
|
||||
Args:
|
||||
image: Input image to draw on (modified in-place).
|
||||
yaw: Yaw angle in degrees.
|
||||
pitch: Pitch angle in degrees.
|
||||
roll: Roll angle in degrees.
|
||||
bbox: Bounding box as ``[x_min, y_min, x_max, y_max]``.
|
||||
size: Cube size in pixels. If None, uses the bounding box width.
|
||||
|
||||
Example:
|
||||
>>> from uniface.draw import draw_head_pose_cube
|
||||
>>> draw_head_pose_cube(image, yaw=10.0, pitch=-5.0, roll=2.0, bbox=[100, 100, 250, 280])
|
||||
"""
|
||||
x_min, y_min, x_max, y_max = map(int, bbox[:4])
|
||||
if size is None:
|
||||
size = x_max - x_min
|
||||
|
||||
h = size * 0.5
|
||||
yaw_r, pitch_r, roll_r = np.radians([-yaw, pitch, roll])
|
||||
|
||||
cx = (x_min + x_max) * 0.5
|
||||
cy = (y_min + y_max) * 0.5
|
||||
|
||||
cos_y, sin_y = np.cos(yaw_r), np.sin(yaw_r)
|
||||
cos_p, sin_p = np.cos(pitch_r), np.sin(pitch_r)
|
||||
cos_r, sin_r = np.cos(roll_r), np.sin(roll_r)
|
||||
|
||||
ex = np.array([cos_y * cos_r, cos_p * sin_r + cos_r * sin_p * sin_y])
|
||||
ey = np.array([-cos_y * sin_r, cos_p * cos_r - sin_p * sin_y * sin_r])
|
||||
ez = np.array([sin_y, -cos_y * sin_p])
|
||||
|
||||
center = np.array([cx, cy])
|
||||
|
||||
def _pt(v: np.ndarray) -> tuple[int, int]:
|
||||
return (int(v[0]), int(v[1]))
|
||||
|
||||
f0 = center + h * (-ex - ey - ez)
|
||||
f1 = center + h * (+ex - ey - ez)
|
||||
f2 = center + h * (+ex + ey - ez)
|
||||
f3 = center + h * (-ex + ey - ez)
|
||||
b0 = center + h * (-ex - ey + ez)
|
||||
b1 = center + h * (+ex - ey + ez)
|
||||
b2 = center + h * (+ex + ey + ez)
|
||||
b3 = center + h * (-ex + ey + ez)
|
||||
|
||||
red = (0, 0, 255)
|
||||
green = (0, 255, 0)
|
||||
blue = (255, 0, 0)
|
||||
|
||||
# Front face at head (red)
|
||||
cv2.line(image, _pt(f0), _pt(f1), red, 2)
|
||||
cv2.line(image, _pt(f1), _pt(f2), red, 2)
|
||||
cv2.line(image, _pt(f2), _pt(f3), red, 2)
|
||||
cv2.line(image, _pt(f3), _pt(f0), red, 2)
|
||||
|
||||
# Back face in looking direction (green)
|
||||
cv2.line(image, _pt(b0), _pt(b1), green, 2)
|
||||
cv2.line(image, _pt(b1), _pt(b2), green, 2)
|
||||
cv2.line(image, _pt(b2), _pt(b3), green, 2)
|
||||
cv2.line(image, _pt(b3), _pt(b0), green, 2)
|
||||
|
||||
# Side edges (blue)
|
||||
cv2.line(image, _pt(f0), _pt(b0), blue, 2)
|
||||
cv2.line(image, _pt(f1), _pt(b1), blue, 2)
|
||||
cv2.line(image, _pt(f2), _pt(b2), blue, 2)
|
||||
cv2.line(image, _pt(f3), _pt(b3), blue, 2)
|
||||
|
||||
|
||||
def draw_head_pose_axis(
|
||||
image: np.ndarray,
|
||||
yaw: float,
|
||||
pitch: float,
|
||||
roll: float,
|
||||
bbox: list[int] | np.ndarray,
|
||||
size_ratio: float = 0.5,
|
||||
) -> None:
|
||||
"""Draw 3D coordinate axes representing head orientation on an image.
|
||||
|
||||
Draws X (red), Y (green), and Z (blue) axes from the center of the
|
||||
bounding box, rotated according to yaw, pitch, and roll.
|
||||
|
||||
Modifies the image in-place.
|
||||
|
||||
Args:
|
||||
image: Input image to draw on (modified in-place).
|
||||
yaw: Yaw angle in degrees.
|
||||
pitch: Pitch angle in degrees.
|
||||
roll: Roll angle in degrees.
|
||||
bbox: Bounding box as ``[x_min, y_min, x_max, y_max]``.
|
||||
size_ratio: Axis length as a fraction of bbox size. Defaults to 0.5.
|
||||
|
||||
Example:
|
||||
>>> from uniface.draw import draw_head_pose_axis
|
||||
>>> draw_head_pose_axis(image, yaw=10.0, pitch=-5.0, roll=2.0, bbox=[100, 100, 250, 280])
|
||||
"""
|
||||
x_min, y_min, x_max, y_max = map(int, bbox[:4])
|
||||
yaw_r, pitch_r, roll_r = np.radians([-yaw, pitch, roll])
|
||||
|
||||
tdx = int(x_min + (x_max - x_min) * 0.5)
|
||||
tdy = int(y_min + (y_max - y_min) * 0.5)
|
||||
|
||||
bbox_size = min(x_max - x_min, y_max - y_min)
|
||||
size = bbox_size * size_ratio
|
||||
|
||||
cos_yaw, sin_yaw = np.cos(yaw_r), np.sin(yaw_r)
|
||||
cos_pitch, sin_pitch = np.cos(pitch_r), np.sin(pitch_r)
|
||||
cos_roll, sin_roll = np.cos(roll_r), np.sin(roll_r)
|
||||
|
||||
# X-Axis (red)
|
||||
x1 = int(size * (cos_yaw * cos_roll) + tdx)
|
||||
y1 = int(size * (cos_pitch * sin_roll + cos_roll * sin_pitch * sin_yaw) + tdy)
|
||||
|
||||
# Y-Axis (green)
|
||||
x2 = int(size * (-cos_yaw * sin_roll) + tdx)
|
||||
y2 = int(size * (cos_pitch * cos_roll - sin_pitch * sin_yaw * sin_roll) + tdy)
|
||||
|
||||
# Z-Axis (blue)
|
||||
x3 = int(size * sin_yaw + tdx)
|
||||
y3 = int(size * (-cos_yaw * sin_pitch) + tdy)
|
||||
|
||||
cv2.line(image, (tdx, tdy), (x1, y1), (0, 0, 255), 2)
|
||||
cv2.line(image, (tdx, tdy), (x2, y2), (0, 255, 0), 2)
|
||||
cv2.line(image, (tdx, tdy), (x3, y3), (255, 0, 0), 2)
|
||||
|
||||
|
||||
def draw_head_pose(
|
||||
image: np.ndarray,
|
||||
bbox: np.ndarray | list[int],
|
||||
pitch: float,
|
||||
yaw: float,
|
||||
roll: float,
|
||||
*,
|
||||
draw_type: str = 'cube',
|
||||
draw_bbox: bool = False,
|
||||
corner_bbox: bool = True,
|
||||
draw_angles: bool = True,
|
||||
) -> None:
|
||||
"""Draw head pose visualization with optional bounding box on an image.
|
||||
|
||||
High-level convenience function that combines bounding box drawing with
|
||||
a 3D shape visualization of head orientation.
|
||||
|
||||
Modifies the image in-place.
|
||||
|
||||
Args:
|
||||
image: Input image to draw on (modified in-place).
|
||||
bbox: Face bounding box in xyxy format ``[x1, y1, x2, y2]``.
|
||||
pitch: Pitch angle in degrees (rotation around X-axis).
|
||||
yaw: Yaw angle in degrees (rotation around Y-axis).
|
||||
roll: Roll angle in degrees (rotation around Z-axis).
|
||||
draw_type: Visualization type, ``'cube'`` or ``'axis'``.
|
||||
Defaults to ``'cube'``.
|
||||
draw_bbox: Whether to draw the bounding box. Defaults to False.
|
||||
corner_bbox: Use corner-style bounding box. Defaults to True.
|
||||
draw_angles: Whether to display angle values as text. Defaults to True.
|
||||
|
||||
Example:
|
||||
>>> from uniface.headpose import HeadPose
|
||||
>>> from uniface.draw import draw_head_pose
|
||||
>>> estimator = HeadPose()
|
||||
>>> result = estimator.estimate(face_crop)
|
||||
>>> draw_head_pose(image, bbox, result.pitch, result.yaw, result.roll)
|
||||
"""
|
||||
x_min, y_min, x_max, y_max = map(int, bbox[:4])
|
||||
|
||||
line_thickness = max(round(sum(image.shape[:2]) / 2 * 0.003), 2)
|
||||
|
||||
if draw_bbox:
|
||||
if corner_bbox:
|
||||
draw_corner_bbox(image, np.array(bbox), color=(0, 255, 0), thickness=line_thickness)
|
||||
else:
|
||||
cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), line_thickness)
|
||||
|
||||
bbox_list = [x_min, y_min, x_max, y_max]
|
||||
if draw_type == 'axis':
|
||||
draw_head_pose_axis(image, yaw, pitch, roll, bbox_list)
|
||||
else:
|
||||
draw_head_pose_cube(image, yaw, pitch, roll, bbox_list)
|
||||
|
||||
if draw_angles:
|
||||
font_scale = max(0.4, min(0.7, (y_max - y_min) / 200))
|
||||
draw_text_label(
|
||||
image,
|
||||
f'P:{pitch:.0f} Y:{yaw:.0f} R:{roll:.0f}',
|
||||
x_min,
|
||||
y_min,
|
||||
bg_color=(0, 0, 255),
|
||||
text_color=(255, 255, 255),
|
||||
font_scale=font_scale,
|
||||
)
|
||||
|
||||
|
||||
def draw_tracks(
|
||||
*,
|
||||
image: np.ndarray,
|
||||
|
||||
@@ -71,8 +71,13 @@ def estimate_norm(
|
||||
alignment[:, 0] += diff_x
|
||||
|
||||
# Compute the transformation matrix
|
||||
transform = SimilarityTransform()
|
||||
transform.estimate(landmark, alignment)
|
||||
try:
|
||||
# scikit-image >= 0.26
|
||||
transform = SimilarityTransform.from_estimate(landmark, alignment)
|
||||
except AttributeError:
|
||||
# scikit-image < 0.26 (e.g. Python 3.10 with older scikit-image)
|
||||
transform = SimilarityTransform()
|
||||
transform.estimate(landmark, alignment)
|
||||
|
||||
matrix = transform.params[0:2, :]
|
||||
inverse_matrix = np.linalg.inv(transform.params)[0:2, :]
|
||||
|
||||
53
uniface/headpose/__init__.py
Normal file
53
uniface/headpose/__init__.py
Normal file
@@ -0,0 +1,53 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
from uniface.types import HeadPoseResult
|
||||
|
||||
from .base import BaseHeadPoseEstimator
|
||||
from .models import HeadPose
|
||||
|
||||
|
||||
def create_head_pose_estimator(method: str = 'headpose', **kwargs) -> BaseHeadPoseEstimator:
|
||||
"""
|
||||
Factory function to create head pose estimators.
|
||||
|
||||
This function initializes and returns a head pose estimator instance based on the
|
||||
specified method. It acts as a high-level interface to the underlying model classes.
|
||||
|
||||
Args:
|
||||
method (str): The head pose estimation method to use.
|
||||
Options: 'headpose' (default).
|
||||
**kwargs: Model-specific parameters passed to the estimator's constructor.
|
||||
For example, `model_name` can be used to select a specific
|
||||
backbone from `HeadPoseWeights` enum (RESNET18, RESNET34, RESNET50,
|
||||
MOBILENET_V2, MOBILENET_V3_SMALL, MOBILENET_V3_LARGE).
|
||||
|
||||
Returns:
|
||||
BaseHeadPoseEstimator: An initialized head pose estimator instance ready for use.
|
||||
|
||||
Raises:
|
||||
ValueError: If the specified `method` is not supported.
|
||||
|
||||
Examples:
|
||||
>>> # Create the default head pose estimator (ResNet18 backbone)
|
||||
>>> estimator = create_head_pose_estimator()
|
||||
|
||||
>>> # Create with MobileNetV2 backbone
|
||||
>>> from uniface.constants import HeadPoseWeights
|
||||
>>> estimator = create_head_pose_estimator('headpose', model_name=HeadPoseWeights.MOBILENET_V2)
|
||||
|
||||
>>> # Use the estimator
|
||||
>>> result = estimator.estimate(face_crop)
|
||||
>>> print(f'Pitch: {result.pitch:.1f}°, Yaw: {result.yaw:.1f}°, Roll: {result.roll:.1f}°')
|
||||
"""
|
||||
method = method.lower()
|
||||
|
||||
if method in ('headpose', 'head_pose', '6drepnet'):
|
||||
return HeadPose(**kwargs)
|
||||
else:
|
||||
available = ['headpose']
|
||||
raise ValueError(f"Unsupported head pose estimation method: '{method}'. Available: {available}")
|
||||
|
||||
|
||||
__all__ = ['BaseHeadPoseEstimator', 'HeadPose', 'HeadPoseResult', 'create_head_pose_estimator']
|
||||
115
uniface/headpose/base.py
Normal file
115
uniface/headpose/base.py
Normal file
@@ -0,0 +1,115 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import numpy as np
|
||||
|
||||
from uniface.types import HeadPoseResult
|
||||
|
||||
__all__ = ['BaseHeadPoseEstimator', 'HeadPoseResult']
|
||||
|
||||
|
||||
class BaseHeadPoseEstimator(ABC):
|
||||
"""
|
||||
Abstract base class for all head pose estimation models.
|
||||
|
||||
This class defines the common interface that all head pose estimators must implement,
|
||||
ensuring consistency across different head pose estimation methods. Head pose estimation
|
||||
predicts the orientation of a person's head based on their face image.
|
||||
|
||||
The head orientation is represented as Euler angles in degrees:
|
||||
- Pitch: Rotation around X-axis (positive = looking down, negative = looking up)
|
||||
- Yaw: Rotation around Y-axis (positive = looking right, negative = looking left)
|
||||
- Roll: Rotation around Z-axis (positive = tilting clockwise, negative = tilting counter-clockwise)
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def _initialize_model(self) -> None:
|
||||
"""
|
||||
Initialize the underlying model for inference.
|
||||
|
||||
This method should handle loading model weights, creating the
|
||||
inference session (e.g., ONNX Runtime), and any necessary
|
||||
setup procedures to prepare the model for prediction.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the model fails to load or initialize.
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the _initialize_model method.')
|
||||
|
||||
@abstractmethod
|
||||
def preprocess(self, face_image: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Preprocess the input face image for model inference.
|
||||
|
||||
This method should take a raw face crop and convert it into the format
|
||||
expected by the model's inference engine (e.g., normalized tensor).
|
||||
|
||||
Args:
|
||||
face_image (np.ndarray): A cropped face image in BGR format with
|
||||
shape (H, W, C).
|
||||
|
||||
Returns:
|
||||
np.ndarray: The preprocessed image tensor ready for inference,
|
||||
typically with shape (1, C, H, W).
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the preprocess method.')
|
||||
|
||||
@abstractmethod
|
||||
def postprocess(self, rotation_matrix: np.ndarray) -> HeadPoseResult:
|
||||
"""
|
||||
Postprocess a rotation matrix into Euler angles.
|
||||
|
||||
This method takes the raw rotation matrix output from the model's
|
||||
inference and converts it into pitch, yaw, and roll angles in degrees.
|
||||
|
||||
Args:
|
||||
rotation_matrix: Rotation matrix with shape (B, 3, 3) from the
|
||||
model inference.
|
||||
|
||||
Returns:
|
||||
HeadPoseResult: Result containing pitch, yaw, and roll in degrees.
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the postprocess method.')
|
||||
|
||||
@abstractmethod
|
||||
def estimate(self, face_image: np.ndarray) -> HeadPoseResult:
|
||||
"""
|
||||
Perform end-to-end head pose estimation on a face image.
|
||||
|
||||
This method orchestrates the full pipeline: preprocessing the input,
|
||||
running inference, and postprocessing to return the head orientation.
|
||||
|
||||
Args:
|
||||
face_image (np.ndarray): A cropped face image in BGR format.
|
||||
The face should be roughly centered and
|
||||
well-framed within the image.
|
||||
|
||||
Returns:
|
||||
HeadPoseResult: Result containing Euler angles in degrees:
|
||||
- pitch: Rotation around X-axis (positive = down)
|
||||
- yaw: Rotation around Y-axis (positive = right)
|
||||
- roll: Rotation around Z-axis (positive = clockwise)
|
||||
|
||||
Example:
|
||||
>>> estimator = create_head_pose_estimator()
|
||||
>>> result = estimator.estimate(face_crop)
|
||||
>>> print(f'Pose: pitch={result.pitch:.1f}°, yaw={result.yaw:.1f}°, roll={result.roll:.1f}°')
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the estimate method.')
|
||||
|
||||
def __call__(self, face_image: np.ndarray) -> HeadPoseResult:
|
||||
"""
|
||||
Provides a convenient, callable shortcut for the `estimate` method.
|
||||
|
||||
Args:
|
||||
face_image (np.ndarray): A cropped face image in BGR format.
|
||||
|
||||
Returns:
|
||||
HeadPoseResult: Result containing pitch, yaw, and roll in degrees.
|
||||
"""
|
||||
return self.estimate(face_image)
|
||||
178
uniface/headpose/models.py
Normal file
178
uniface/headpose/models.py
Normal file
@@ -0,0 +1,178 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from uniface.constants import HeadPoseWeights
|
||||
from uniface.log import Logger
|
||||
from uniface.model_store import verify_model_weights
|
||||
from uniface.onnx_utils import create_onnx_session
|
||||
from uniface.types import HeadPoseResult
|
||||
|
||||
from .base import BaseHeadPoseEstimator
|
||||
|
||||
__all__ = ['HeadPose']
|
||||
|
||||
|
||||
class HeadPose(BaseHeadPoseEstimator):
|
||||
"""
|
||||
Head Pose Estimation with ONNX Runtime using 6D Rotation Representation.
|
||||
|
||||
This model estimates head orientation from a single face image by predicting
|
||||
a 3x3 rotation matrix (via continuous 6D representation) and converting it to
|
||||
Euler angles (pitch, yaw, roll) in degrees.
|
||||
|
||||
Supports multiple backbone architectures: ResNet-18/34/50, MobileNetV2,
|
||||
and MobileNetV3 (small/large).
|
||||
|
||||
Reference:
|
||||
https://github.com/yakhyo/head-pose-estimation
|
||||
|
||||
Args:
|
||||
model_name (HeadPoseWeights): The enum specifying the head pose model to load.
|
||||
Options: RESNET18, RESNET34, RESNET50, MOBILENET_V2, MOBILENET_V3_SMALL,
|
||||
MOBILENET_V3_LARGE. Defaults to `HeadPoseWeights.RESNET18`.
|
||||
input_size (tuple[int, int]): The resolution (width, height) for the model's
|
||||
input. Defaults to (224, 224).
|
||||
providers (list[str] | None): ONNX Runtime execution providers. If None, auto-detects
|
||||
the best available provider. Example: ['CPUExecutionProvider'] to force CPU.
|
||||
|
||||
Attributes:
|
||||
input_size (tuple[int, int]): Model input dimensions.
|
||||
input_mean (np.ndarray): Per-channel mean values for normalization (ImageNet).
|
||||
input_std (np.ndarray): Per-channel std values for normalization (ImageNet).
|
||||
|
||||
Example:
|
||||
>>> from uniface.headpose import HeadPose
|
||||
>>> from uniface import RetinaFace
|
||||
>>>
|
||||
>>> detector = RetinaFace()
|
||||
>>> head_pose = HeadPose()
|
||||
>>>
|
||||
>>> faces = detector.detect(image)
|
||||
>>> for face in faces:
|
||||
... bbox = face.bbox
|
||||
... x1, y1, x2, y2 = map(int, bbox[:4])
|
||||
... face_crop = image[y1:y2, x1:x2]
|
||||
... result = head_pose.estimate(face_crop)
|
||||
... print(f'Pose: pitch={result.pitch:.1f}°, yaw={result.yaw:.1f}°, roll={result.roll:.1f}°')
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: HeadPoseWeights = HeadPoseWeights.RESNET18,
|
||||
input_size: tuple[int, int] = (224, 224),
|
||||
providers: list[str] | None = None,
|
||||
) -> None:
|
||||
Logger.info(f'Initializing HeadPose with model={model_name}, input_size={input_size}')
|
||||
|
||||
self.input_size = input_size
|
||||
self.input_mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
self.input_std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
self.providers = providers
|
||||
|
||||
self.model_path = verify_model_weights(model_name)
|
||||
self._initialize_model()
|
||||
|
||||
def _initialize_model(self) -> None:
|
||||
"""
|
||||
Initialize the ONNX model from the stored model path.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the model fails to load or initialize.
|
||||
"""
|
||||
try:
|
||||
self.session = create_onnx_session(self.model_path, providers=self.providers)
|
||||
|
||||
input_cfg = self.session.get_inputs()[0]
|
||||
input_shape = input_cfg.shape
|
||||
self.input_name = input_cfg.name
|
||||
self.input_size = tuple(input_shape[2:4][::-1])
|
||||
|
||||
outputs = self.session.get_outputs()
|
||||
self.output_names = [output.name for output in outputs]
|
||||
|
||||
if len(self.output_names) != 1:
|
||||
raise ValueError(f'Expected 1 output node (rotation_matrix), got {len(self.output_names)}')
|
||||
|
||||
Logger.info(f'HeadPose initialized with input size {self.input_size}')
|
||||
|
||||
except Exception as e:
|
||||
Logger.error(f"Failed to load head pose model from '{self.model_path}'", exc_info=True)
|
||||
raise RuntimeError(f'Failed to initialize head pose model: {e}') from e
|
||||
|
||||
def preprocess(self, face_image: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Preprocess a face crop for head pose estimation.
|
||||
|
||||
Args:
|
||||
face_image (np.ndarray): A cropped face image in BGR format.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Preprocessed image tensor with shape (1, 3, H, W).
|
||||
"""
|
||||
image = cv2.cvtColor(face_image, cv2.COLOR_BGR2RGB)
|
||||
image = cv2.resize(image, self.input_size)
|
||||
image = image.astype(np.float32) / 255.0
|
||||
image = (image - self.input_mean) / self.input_std
|
||||
|
||||
# HWC -> CHW -> NCHW
|
||||
image = np.transpose(image, (2, 0, 1))
|
||||
image = np.expand_dims(image, axis=0).astype(np.float32)
|
||||
|
||||
return image
|
||||
|
||||
@staticmethod
|
||||
def rotation_matrix_to_euler(rotation_matrix: np.ndarray) -> np.ndarray:
|
||||
"""Convert (B, 3, 3) rotation matrices to Euler angles in degrees.
|
||||
|
||||
Uses the ZYX convention to decompose rotation matrices into
|
||||
pitch (X), yaw (Y), and roll (Z) angles.
|
||||
|
||||
Args:
|
||||
rotation_matrix: Batch of rotation matrices with shape (B, 3, 3).
|
||||
|
||||
Returns:
|
||||
np.ndarray: Euler angles with shape (B, 3) as [pitch, yaw, roll] in degrees.
|
||||
"""
|
||||
R = rotation_matrix
|
||||
sy = np.sqrt(R[:, 0, 0] ** 2 + R[:, 1, 0] ** 2)
|
||||
singular = sy < 1e-6
|
||||
|
||||
x = np.where(singular, np.arctan2(-R[:, 1, 2], R[:, 1, 1]), np.arctan2(R[:, 2, 1], R[:, 2, 2]))
|
||||
y = np.arctan2(-R[:, 2, 0], sy)
|
||||
z = np.where(singular, np.zeros_like(sy), np.arctan2(R[:, 1, 0], R[:, 0, 0]))
|
||||
|
||||
return np.degrees(np.stack([x, y, z], axis=1))
|
||||
|
||||
def postprocess(self, rotation_matrix: np.ndarray) -> HeadPoseResult:
|
||||
"""
|
||||
Convert a rotation matrix into Euler angles.
|
||||
|
||||
Args:
|
||||
rotation_matrix: Rotation matrix with shape (B, 3, 3).
|
||||
|
||||
Returns:
|
||||
HeadPoseResult: Result containing pitch, yaw, and roll in degrees.
|
||||
"""
|
||||
euler = self.rotation_matrix_to_euler(rotation_matrix)
|
||||
return HeadPoseResult(
|
||||
pitch=float(euler[0, 0]),
|
||||
yaw=float(euler[0, 1]),
|
||||
roll=float(euler[0, 2]),
|
||||
)
|
||||
|
||||
def estimate(self, face_image: np.ndarray) -> HeadPoseResult:
|
||||
"""
|
||||
Perform end-to-end head pose estimation on a face image.
|
||||
|
||||
This method orchestrates the full pipeline: preprocessing the input,
|
||||
running inference, and postprocessing to return the head orientation.
|
||||
"""
|
||||
input_tensor = self.preprocess(face_image)
|
||||
outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
|
||||
rotation_matrix = outputs[0] # (1, 3, 3)
|
||||
|
||||
return self.postprocess(rotation_matrix)
|
||||
@@ -1,9 +0,0 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
"""Vector indexing backends for fast similarity search."""
|
||||
|
||||
from uniface.indexing.faiss import FAISS
|
||||
|
||||
__all__ = ['FAISS']
|
||||
@@ -10,17 +10,23 @@ import numpy as np
|
||||
|
||||
|
||||
class BaseFaceParser(ABC):
|
||||
"""
|
||||
Abstract base class for all face parsing models.
|
||||
"""Abstract base class for all face parsing models.
|
||||
|
||||
This class defines the common interface that all face parsing models must implement,
|
||||
ensuring consistency across different parsing methods. Face parsing segments a face
|
||||
image into semantic regions such as skin, eyes, nose, mouth, hair, etc.
|
||||
|
||||
The output is a segmentation mask where each pixel is assigned a class label
|
||||
representing a facial component.
|
||||
Subclasses must define a ``mask_type`` class attribute to indicate output format:
|
||||
|
||||
- ``"class_ids"``: uint8 mask with discrete class labels (e.g. BiSeNet: 0-18)
|
||||
- ``"probability"``: float32 mask with continuous values in [0, 1] (e.g. XSeg)
|
||||
|
||||
Attributes:
|
||||
mask_type (str): Output format identifier. Must be set by subclasses.
|
||||
"""
|
||||
|
||||
mask_type: str
|
||||
|
||||
@abstractmethod
|
||||
def _initialize_model(self) -> None:
|
||||
"""
|
||||
@@ -86,13 +92,17 @@ class BaseFaceParser(ABC):
|
||||
Ignored by parsers that do not need landmarks (e.g., BiSeNet).
|
||||
|
||||
Returns:
|
||||
np.ndarray: Segmentation mask with the same size as input image,
|
||||
where each pixel value represents a facial component class.
|
||||
np.ndarray: Segmentation mask with the same size as input image.
|
||||
Format depends on ``mask_type``:
|
||||
|
||||
- ``"class_ids"``: uint8 with discrete class labels
|
||||
- ``"probability"``: float32 with values in [0, 1]
|
||||
|
||||
Example:
|
||||
>>> parser = create_face_parser()
|
||||
>>> mask = parser.parse(face_crop)
|
||||
>>> print(f'Mask shape: {mask.shape}, unique classes: {np.unique(mask)}')
|
||||
>>> print(f'Mask type: {parser.mask_type}')
|
||||
>>> print(f'Mask shape: {mask.shape}, dtype: {mask.dtype}')
|
||||
"""
|
||||
raise NotImplementedError('Subclasses must implement the parse method.')
|
||||
|
||||
|
||||
@@ -18,8 +18,7 @@ __all__ = ['BiSeNet']
|
||||
|
||||
|
||||
class BiSeNet(BaseFaceParser):
|
||||
"""
|
||||
BiSeNet: Bilateral Segmentation Network for Face Parsing with ONNX Runtime.
|
||||
"""BiSeNet: Bilateral Segmentation Network for Face Parsing with ONNX Runtime.
|
||||
|
||||
BiSeNet is a semantic segmentation model that segments a face image into
|
||||
different facial components such as skin, eyes, nose, mouth, hair, etc. The model
|
||||
@@ -45,6 +44,7 @@ class BiSeNet(BaseFaceParser):
|
||||
input_size (Tuple[int, int]): Model input dimensions.
|
||||
input_mean (np.ndarray): Per-channel mean values for normalization (ImageNet).
|
||||
input_std (np.ndarray): Per-channel std values for normalization (ImageNet).
|
||||
mask_type (str): Output type identifier - "class_ids" for BiSeNet.
|
||||
|
||||
Example:
|
||||
>>> from uniface.parsing import BiSeNet
|
||||
@@ -61,8 +61,11 @@ class BiSeNet(BaseFaceParser):
|
||||
... face_crop = image[y1:y2, x1:x2]
|
||||
... mask = parser.parse(face_crop)
|
||||
... print(f'Mask shape: {mask.shape}, unique classes: {np.unique(mask)}')
|
||||
... print(f'Output type: {parser.mask_type}') # "class_ids"
|
||||
"""
|
||||
|
||||
mask_type = 'class_ids'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: ParsingWeights = ParsingWeights.RESNET18,
|
||||
|
||||
@@ -19,10 +19,9 @@ __all__ = ['XSeg']
|
||||
|
||||
|
||||
class XSeg(BaseFaceParser):
|
||||
"""
|
||||
XSeg: Face Segmentation Model from DeepFaceLab with ONNX Runtime.
|
||||
"""XSeg: Face Segmentation Model from DeepFaceLab with ONNX Runtime.
|
||||
|
||||
XSeg outputs a mask for face regions. Unlike BiSeNet which works
|
||||
XSeg outputs a soft probability mask for face regions. Unlike BiSeNet which works
|
||||
on bbox crops, XSeg requires 5-point landmarks for face alignment. The model
|
||||
uses NHWC input format and outputs values in [0, 1] range.
|
||||
|
||||
@@ -43,6 +42,7 @@ class XSeg(BaseFaceParser):
|
||||
align_size (int): Face alignment output size.
|
||||
blur_sigma (float): Blur sigma for post-processing.
|
||||
input_size (tuple[int, int]): Model input dimensions (width, height).
|
||||
mask_type (str): Output type identifier - "probability" for XSeg.
|
||||
|
||||
Example:
|
||||
>>> from uniface.parsing import XSeg
|
||||
@@ -56,8 +56,11 @@ class XSeg(BaseFaceParser):
|
||||
... if face.landmarks is not None:
|
||||
... mask = parser.parse(image, landmarks=face.landmarks)
|
||||
... print(f'Mask shape: {mask.shape}')
|
||||
... print(f'Output type: {parser.mask_type}') # "probability"
|
||||
"""
|
||||
|
||||
mask_type = 'probability'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: XSegWeights = XSegWeights.DEFAULT,
|
||||
|
||||
@@ -141,7 +141,7 @@ class BaseRecognizer(ABC):
|
||||
image is already aligned.
|
||||
|
||||
Returns:
|
||||
Face embedding vector (typically 512-dimensional).
|
||||
Face embedding with shape (1, 512) — raw ONNX output with batch dimension.
|
||||
"""
|
||||
# If landmarks are provided, align the face first
|
||||
if landmarks is not None:
|
||||
@@ -164,9 +164,9 @@ class BaseRecognizer(ABC):
|
||||
landmarks: Facial landmarks (5 points for alignment).
|
||||
|
||||
Returns:
|
||||
L2-normalized face embedding vector (typically 512-dimensional).
|
||||
L2-normalized face embedding as a 1D vector with shape (512,).
|
||||
"""
|
||||
embedding = self.get_embedding(image, landmarks)
|
||||
embedding = self.get_embedding(image, landmarks).ravel()
|
||||
norm = np.linalg.norm(embedding)
|
||||
return embedding / norm if norm > 0 else embedding
|
||||
|
||||
@@ -178,6 +178,6 @@ class BaseRecognizer(ABC):
|
||||
landmarks: Facial landmarks (5 points for alignment).
|
||||
|
||||
Returns:
|
||||
L2-normalized face embedding vector (typically 512-dimensional).
|
||||
L2-normalized face embedding as a 1D vector with shape (512,).
|
||||
"""
|
||||
return self.get_normalized_embedding(image, landmarks)
|
||||
|
||||
10
uniface/stores/__init__.py
Normal file
10
uniface/stores/__init__.py
Normal file
@@ -0,0 +1,10 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
"""Vector store backends for fast face embedding similarity search."""
|
||||
|
||||
from uniface.stores.base import BaseStore
|
||||
from uniface.stores.faiss import FAISS
|
||||
|
||||
__all__ = ['BaseStore', 'FAISS']
|
||||
69
uniface/stores/base.py
Normal file
69
uniface/stores/base.py
Normal file
@@ -0,0 +1,69 @@
|
||||
# Copyright 2025-2026 Yakhyokhuja Valikhujaev
|
||||
# Author: Yakhyokhuja Valikhujaev
|
||||
# GitHub: https://github.com/yakhyo
|
||||
|
||||
"""Abstract base class for vector store backends."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
__all__ = ['BaseStore']
|
||||
|
||||
Metadata = dict[str, Any]
|
||||
|
||||
|
||||
class BaseStore(ABC):
|
||||
"""Abstract interface for face embedding vector stores.
|
||||
|
||||
All vector store backends (FAISS, Qdrant, etc.) must implement
|
||||
this interface to ensure consistent usage across the library.
|
||||
|
||||
Embeddings are expected to be L2-normalised so that inner product
|
||||
equals cosine similarity.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def add(self, embedding: np.ndarray, metadata: Metadata) -> None:
|
||||
"""Add a single embedding with associated metadata.
|
||||
|
||||
Args:
|
||||
embedding: L2-normalised embedding vector.
|
||||
metadata: Arbitrary dict of JSON-serialisable key-value pairs.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def search(
|
||||
self,
|
||||
embedding: np.ndarray,
|
||||
threshold: float = 0.4,
|
||||
) -> tuple[Metadata | None, float]:
|
||||
"""Find the closest match for a query embedding.
|
||||
|
||||
Args:
|
||||
embedding: L2-normalised query vector.
|
||||
threshold: Minimum similarity to accept a match.
|
||||
|
||||
Returns:
|
||||
``(metadata, similarity)`` for the best match, or
|
||||
``(None, similarity)`` when below *threshold* or empty.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def remove(self, key: str, value: Any) -> int:
|
||||
"""Remove all entries where ``metadata[key] == value``.
|
||||
|
||||
Args:
|
||||
key: Metadata key to match against.
|
||||
value: Value to match.
|
||||
|
||||
Returns:
|
||||
Number of entries removed.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def __len__(self) -> int:
|
||||
"""Return the number of vectors in the store."""
|
||||
@@ -12,9 +12,9 @@ import numpy as np
|
||||
|
||||
from uniface.log import Logger
|
||||
|
||||
__all__ = ['FAISS']
|
||||
from .base import BaseStore, Metadata
|
||||
|
||||
Metadata = dict[str, Any]
|
||||
__all__ = ['FAISS']
|
||||
|
||||
|
||||
def _import_faiss():
|
||||
@@ -34,7 +34,7 @@ def _import_faiss():
|
||||
return faiss
|
||||
|
||||
|
||||
class FAISS:
|
||||
class FAISS(BaseStore):
|
||||
"""FAISS vector store using IndexFlatIP (inner product).
|
||||
|
||||
Vectors must be L2-normalised **before** being added so that inner
|
||||
@@ -49,7 +49,7 @@ class FAISS:
|
||||
db_path: Directory for persisting the index and metadata.
|
||||
|
||||
Example:
|
||||
>>> from uniface.indexing import FAISS
|
||||
>>> from uniface.stores import FAISS
|
||||
>>> store = FAISS(embedding_size=512, db_path='./my_index')
|
||||
>>> store.add(embedding, {'person_id': '001', 'name': 'Alice'})
|
||||
>>> result, score = store.search(query_embedding)
|
||||
@@ -28,6 +28,7 @@ __all__ = [
|
||||
'EmotionResult',
|
||||
'Face',
|
||||
'GazeResult',
|
||||
'HeadPoseResult',
|
||||
'SpoofingResult',
|
||||
]
|
||||
|
||||
@@ -48,6 +49,24 @@ class GazeResult:
|
||||
return f'GazeResult(pitch={self.pitch:.4f}, yaw={self.yaw:.4f})'
|
||||
|
||||
|
||||
@dataclass(slots=True, frozen=True)
|
||||
class HeadPoseResult:
|
||||
"""Result of head pose estimation.
|
||||
|
||||
Attributes:
|
||||
pitch: Rotation around X-axis in degrees (positive = looking down).
|
||||
yaw: Rotation around Y-axis in degrees (positive = looking right).
|
||||
roll: Rotation around Z-axis in degrees (positive = tilting clockwise).
|
||||
"""
|
||||
|
||||
pitch: float
|
||||
yaw: float
|
||||
roll: float
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f'HeadPoseResult(pitch={self.pitch:.1f}, yaw={self.yaw:.1f}, roll={self.roll:.1f})'
|
||||
|
||||
|
||||
@dataclass(slots=True, frozen=True)
|
||||
class SpoofingResult:
|
||||
"""Result of face anti-spoofing detection.
|
||||
@@ -245,5 +264,5 @@ class Face:
|
||||
if self.emotion is not None:
|
||||
parts.append(f'emotion={self.emotion}')
|
||||
if self.embedding is not None:
|
||||
parts.append(f'embedding_dim={self.embedding.shape[0]}')
|
||||
parts.append(f'embedding_dim={self.embedding.shape[-1]}')
|
||||
return ', '.join(parts) + ')'
|
||||
|
||||
Reference in New Issue
Block a user