Initial commit: scene-actor-extraction pipeline
Source (KPN++ pipeline nodes, ArcFace embedders, SCRFD/YuNet detectors, gallery builder), build scripts, and eval artifacts. - external/KPN as a git submodule (gitea.tourolle.paris/dtourolle/KPN) - ONNX models tracked via Git LFS (models/*.onnx) - generated outputs, TensorRT engines, reference repos, and media ignored
This commit is contained in:
@@ -0,0 +1 @@
|
|||||||
|
models/*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||||
+67
@@ -0,0 +1,67 @@
|
|||||||
|
# Build
|
||||||
|
build/
|
||||||
|
cmake-build-*/
|
||||||
|
CMakeCache.txt
|
||||||
|
CMakeFiles/
|
||||||
|
*.cmake
|
||||||
|
Makefile
|
||||||
|
install_manifest.txt
|
||||||
|
compile_commands.json
|
||||||
|
|
||||||
|
# Compiled objects
|
||||||
|
*.o
|
||||||
|
*.a
|
||||||
|
*.so
|
||||||
|
*.dylib
|
||||||
|
|
||||||
|
# Video files
|
||||||
|
*.mp4
|
||||||
|
*.mkv
|
||||||
|
*.avi
|
||||||
|
*.mov
|
||||||
|
|
||||||
|
# ONNX models in models/ are tracked via Git LFS (see .gitattributes).
|
||||||
|
# Any stray ONNX elsewhere is generated/downloaded and not tracked.
|
||||||
|
external/*.onnx
|
||||||
|
|
||||||
|
# Generated TensorRT engines (rebuilt by ORT / scripts/build_trt_engines.sh)
|
||||||
|
trt_cache/
|
||||||
|
|
||||||
|
# Gallery JSON files (generated)
|
||||||
|
gallery.json
|
||||||
|
gallery_*.json
|
||||||
|
|
||||||
|
# Per-movie analysis output (generated by scene_analyze)
|
||||||
|
death_of_stalin.json
|
||||||
|
|
||||||
|
# Per-frame debug images
|
||||||
|
images/
|
||||||
|
|
||||||
|
# Annotations output
|
||||||
|
annotations.json
|
||||||
|
*_annotations.json
|
||||||
|
|
||||||
|
# MovieNet evaluation data
|
||||||
|
movienet-ps/
|
||||||
|
|
||||||
|
# Eval probe images (data, regenerable)
|
||||||
|
eval/probe/
|
||||||
|
|
||||||
|
# Local reference repos kept for inspiration (each has its own .git)
|
||||||
|
inspiration/
|
||||||
|
|
||||||
|
# Python
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
.venv/
|
||||||
|
venv/
|
||||||
|
|
||||||
|
# Editor / OS
|
||||||
|
.vscode/
|
||||||
|
.idea/
|
||||||
|
.claude/settings.local.json
|
||||||
|
*.swp
|
||||||
|
*.swo
|
||||||
|
.DS_Store
|
||||||
|
Thumbs.db
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
[submodule "external/KPN"]
|
||||||
|
path = external/KPN
|
||||||
|
url = git@gitea.tourolle.paris:dtourolle/KPN.git
|
||||||
+148
@@ -0,0 +1,148 @@
|
|||||||
|
cmake_minimum_required(VERSION 3.21)
|
||||||
|
project(scene_actor_extraction VERSION 0.1.0 LANGUAGES CXX)
|
||||||
|
|
||||||
|
set(CMAKE_CXX_STANDARD 20)
|
||||||
|
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||||
|
set(CMAKE_CXX_EXTENSIONS OFF)
|
||||||
|
|
||||||
|
# ── Dependencies ──────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
# KPN++ (pipeline backbone)
|
||||||
|
set(KPN_BUILD_TESTS OFF CACHE BOOL "" FORCE)
|
||||||
|
set(KPN_BUILD_PYTHON OFF CACHE BOOL "" FORCE)
|
||||||
|
set(KPN_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
|
||||||
|
option(SAE_WEB_DEBUG "Enable KPN web debug UI (localhost:9090)" OFF)
|
||||||
|
if(SAE_WEB_DEBUG)
|
||||||
|
set(KPN_WEB_DEBUG ON CACHE BOOL "" FORCE)
|
||||||
|
endif()
|
||||||
|
add_subdirectory(external/KPN)
|
||||||
|
|
||||||
|
# OpenCV (video decode, image ops, DNN inference, face detection)
|
||||||
|
find_package(OpenCV 4 REQUIRED COMPONENTS
|
||||||
|
core imgproc imgcodecs videoio dnn objdetect highgui)
|
||||||
|
|
||||||
|
# ONNX Runtime (SCRFD detector — cv::dnn cannot handle dynamic Shape nodes)
|
||||||
|
find_library(ORT_LIB onnxruntime REQUIRED
|
||||||
|
HINTS /usr/lib /usr/local/lib)
|
||||||
|
find_path(ORT_INCLUDE onnxruntime_cxx_api.h
|
||||||
|
PATH_SUFFIXES onnxruntime
|
||||||
|
HINTS /usr/include /usr/local/include
|
||||||
|
REQUIRED)
|
||||||
|
add_library(onnxruntime UNKNOWN IMPORTED)
|
||||||
|
set_target_properties(onnxruntime PROPERTIES
|
||||||
|
IMPORTED_LOCATION "${ORT_LIB}"
|
||||||
|
INTERFACE_INCLUDE_DIRECTORIES "${ORT_INCLUDE}")
|
||||||
|
message(STATUS "ONNX Runtime: ${ORT_LIB} headers: ${ORT_INCLUDE}")
|
||||||
|
|
||||||
|
# TensorRT + CUDA runtime (raw-TRT ArcFace embedder; activated by --arcface-engine).
|
||||||
|
find_library(NVINFER_LIB nvinfer
|
||||||
|
HINTS /usr/lib /usr/local/lib /opt/tensorrt/lib)
|
||||||
|
find_path(NVINFER_INCLUDE NvInfer.h
|
||||||
|
HINTS /usr/include /usr/local/include /opt/tensorrt/include)
|
||||||
|
find_library(CUDART_LIB cudart
|
||||||
|
HINTS /opt/cuda/lib64 /usr/local/cuda/lib64 /usr/lib)
|
||||||
|
find_path(CUDART_INCLUDE cuda_runtime_api.h
|
||||||
|
HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include
|
||||||
|
/usr/local/cuda/include /usr/include)
|
||||||
|
if(NOT (NVINFER_LIB AND NVINFER_INCLUDE AND CUDART_LIB AND CUDART_INCLUDE))
|
||||||
|
message(FATAL_ERROR
|
||||||
|
"TensorRT or CUDA runtime not found "
|
||||||
|
"(nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE} "
|
||||||
|
"cudart=${CUDART_LIB} headers=${CUDART_INCLUDE})")
|
||||||
|
endif()
|
||||||
|
add_library(trt_runtime INTERFACE)
|
||||||
|
target_include_directories(trt_runtime INTERFACE
|
||||||
|
"${NVINFER_INCLUDE}" "${CUDART_INCLUDE}")
|
||||||
|
target_link_libraries(trt_runtime INTERFACE
|
||||||
|
"${NVINFER_LIB}" "${CUDART_LIB}")
|
||||||
|
message(STATUS "TensorRT: ${NVINFER_LIB} CUDA runtime: ${CUDART_LIB}")
|
||||||
|
|
||||||
|
# FFmpeg (NVDEC hardware video decode + swscale colour conversion)
|
||||||
|
find_package(PkgConfig REQUIRED)
|
||||||
|
pkg_check_modules(AVFORMAT REQUIRED libavformat)
|
||||||
|
pkg_check_modules(AVCODEC REQUIRED libavcodec)
|
||||||
|
pkg_check_modules(AVUTIL REQUIRED libavutil)
|
||||||
|
pkg_check_modules(SWSCALE REQUIRED libswscale)
|
||||||
|
|
||||||
|
add_library(ffmpeg_libs INTERFACE)
|
||||||
|
target_compile_options(ffmpeg_libs INTERFACE
|
||||||
|
${AVFORMAT_CFLAGS_OTHER} ${AVCODEC_CFLAGS_OTHER}
|
||||||
|
${AVUTIL_CFLAGS_OTHER} ${SWSCALE_CFLAGS_OTHER})
|
||||||
|
target_include_directories(ffmpeg_libs INTERFACE
|
||||||
|
${AVFORMAT_INCLUDE_DIRS} ${AVCODEC_INCLUDE_DIRS}
|
||||||
|
${AVUTIL_INCLUDE_DIRS} ${SWSCALE_INCLUDE_DIRS})
|
||||||
|
target_link_libraries(ffmpeg_libs INTERFACE
|
||||||
|
${AVFORMAT_LIBRARIES} ${AVCODEC_LIBRARIES}
|
||||||
|
${AVUTIL_LIBRARIES} ${SWSCALE_LIBRARIES})
|
||||||
|
message(STATUS "FFmpeg: avformat=${AVFORMAT_VERSION} avcodec=${AVCODEC_VERSION}")
|
||||||
|
|
||||||
|
# nlohmann/json (gallery + output serialisation)
|
||||||
|
include(FetchContent)
|
||||||
|
FetchContent_Declare(
|
||||||
|
nlohmann_json
|
||||||
|
GIT_REPOSITORY https://github.com/nlohmann/json.git
|
||||||
|
GIT_TAG v3.11.3
|
||||||
|
GIT_SHALLOW TRUE
|
||||||
|
)
|
||||||
|
FetchContent_MakeAvailable(nlohmann_json)
|
||||||
|
|
||||||
|
# ── Model paths ───────────────────────────────────────────────────────────────
|
||||||
|
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
|
||||||
|
CACHE PATH "Directory containing ONNX model files")
|
||||||
|
|
||||||
|
# ── Shared library: gallery store ─────────────────────────────────────────────
|
||||||
|
add_library(sae_gallery STATIC
|
||||||
|
src/gallery/gallery_store.cpp
|
||||||
|
src/gallery/gallery_builder.cpp
|
||||||
|
)
|
||||||
|
target_include_directories(sae_gallery PUBLIC src)
|
||||||
|
target_link_libraries(sae_gallery PUBLIC
|
||||||
|
kpn
|
||||||
|
${OpenCV_LIBS}
|
||||||
|
nlohmann_json::nlohmann_json
|
||||||
|
onnxruntime
|
||||||
|
trt_runtime
|
||||||
|
ffmpeg_libs
|
||||||
|
)
|
||||||
|
target_compile_definitions(sae_gallery PUBLIC
|
||||||
|
SAE_MODELS_DIR="${SAE_MODELS_DIR}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── embed_faces — image → embedding JSON (used by gallery builder scripts) ────
|
||||||
|
add_executable(embed_faces src/embed_faces.cpp)
|
||||||
|
target_link_libraries(embed_faces PRIVATE
|
||||||
|
kpn
|
||||||
|
${OpenCV_LIBS}
|
||||||
|
nlohmann_json::nlohmann_json
|
||||||
|
onnxruntime
|
||||||
|
trt_runtime
|
||||||
|
ffmpeg_libs
|
||||||
|
)
|
||||||
|
target_compile_definitions(embed_faces PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")
|
||||||
|
|
||||||
|
# ── analyze — main analysis binary ───────────────────────────────────────────
|
||||||
|
add_executable(scene_analyze src/main.cpp)
|
||||||
|
target_link_libraries(scene_analyze PRIVATE sae_gallery)
|
||||||
|
|
||||||
|
# ── analyze_debug — same binary with debug frame/crop output ─────────────────
|
||||||
|
add_executable(scene_analyze_debug src/main.cpp)
|
||||||
|
target_link_libraries(scene_analyze_debug PRIVATE sae_gallery)
|
||||||
|
target_compile_definitions(scene_analyze_debug PRIVATE SAE_DEBUG=1)
|
||||||
|
|
||||||
|
# ── scene_preview — live annotated display while analysing ───────────────────
|
||||||
|
add_executable(scene_preview src/scene_preview.cpp)
|
||||||
|
target_link_libraries(scene_preview PRIVATE sae_gallery)
|
||||||
|
|
||||||
|
# ── build_gallery — offline gallery construction tool ────────────────────────
|
||||||
|
add_executable(build_gallery src/build_gallery.cpp)
|
||||||
|
target_link_libraries(build_gallery PRIVATE sae_gallery)
|
||||||
|
|
||||||
|
# ── Optional: web debug UI for pipeline introspection ────────────────────────
|
||||||
|
if(SAE_WEB_DEBUG)
|
||||||
|
kpn_target_enable_web_debug(scene_analyze)
|
||||||
|
kpn_target_enable_web_debug(scene_analyze_debug)
|
||||||
|
kpn_target_enable_web_debug(scene_preview)
|
||||||
|
endif()
|
||||||
|
|
||||||
|
message(STATUS "OpenCV ${OpenCV_VERSION} found")
|
||||||
|
message(STATUS "Models dir: ${SAE_MODELS_DIR}")
|
||||||
@@ -0,0 +1,107 @@
|
|||||||
|
# Scene Actor Extraction
|
||||||
|
|
||||||
|
Identifies actors in movie files and produces X-ray-style scene annotations compatible with [Jellyfin](https://jellyfin.org/). Built on a KPN++ pipeline with ArcFace embeddings and a tracked-identity matcher.
|
||||||
|
|
||||||
|
## How it works
|
||||||
|
|
||||||
|
1. **Build a gallery** — download actor headshots from TMDB/IMDB, embed them with ArcFace (`build_gallery` / `scripts/make_gallery.py`).
|
||||||
|
2. **Analyze a movie** — `scene_analyze` decodes frames at configurable FPS, detects faces (YuNet/SCRFD), tracks them across cuts, matches identities against the gallery using calibrated similarity, and writes time-window JSON.
|
||||||
|
3. **Output** — minimal mode produces Jellyfin-ready actor name + time-window JSON; standard mode adds per-frame bbox, similarity, and track data.
|
||||||
|
|
||||||
|
## Dependencies
|
||||||
|
|
||||||
|
| Dependency | Role |
|
||||||
|
|---|---|
|
||||||
|
| KPN++ | Pipeline backbone (nodes, networks) |
|
||||||
|
| OpenCV 4 | Video decode, image ops, DNN inference, YuNet face detection |
|
||||||
|
| ONNX Runtime | SCRFD face detector (dynamic shape nodes unsupported by cv::dnn) |
|
||||||
|
| nlohmann/json | JSON I/O |
|
||||||
|
|
||||||
|
## Build
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cmake -B build -DCMAKE_BUILD_TYPE=Release
|
||||||
|
cmake --build build -j$(nproc)
|
||||||
|
```
|
||||||
|
|
||||||
|
Optional flags:
|
||||||
|
|
||||||
|
| Flag | Default | Effect |
|
||||||
|
|---|---|---|
|
||||||
|
| `-DSAE_WEB_DEBUG=ON` | OFF | Enables KPN web debug UI at `localhost:9090` |
|
||||||
|
|
||||||
|
## Models
|
||||||
|
|
||||||
|
Download the required ONNX models:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
bash scripts/download_models.sh
|
||||||
|
```
|
||||||
|
|
||||||
|
Models are placed in `external/`:
|
||||||
|
- `arcface_w600k_r50.onnx` — primary ArcFace embedder
|
||||||
|
- `arcface_w600k_mbf.onnx`, `arcface_r18.onnx` — lighter alternatives
|
||||||
|
- `face_detection_yunet_2023mar.onnx` — YuNet face detector
|
||||||
|
- `scrfd_500m_bnkps.onnx` — SCRFD face detector
|
||||||
|
|
||||||
|
## Binaries
|
||||||
|
|
||||||
|
| Binary | Description |
|
||||||
|
|---|---|
|
||||||
|
| `scene_analyze` | Main analysis pipeline, writes JSON output |
|
||||||
|
| `scene_analyze_debug` | Same as above + per-frame annotated JPEGs (`SAE_DEBUG=1`) |
|
||||||
|
| `scene_preview` | Live OpenCV display window while analysing |
|
||||||
|
| `build_gallery` | Offline gallery builder from a directory of images |
|
||||||
|
| `embed_faces` | Standalone embedder used by gallery scripts |
|
||||||
|
|
||||||
|
### `scene_analyze`
|
||||||
|
|
||||||
|
```bash
|
||||||
|
./build/scene_analyze --gallery gallery.json --input movie.mp4 [options]
|
||||||
|
```
|
||||||
|
|
||||||
|
Key options:
|
||||||
|
|
||||||
|
| Flag | Default | Description |
|
||||||
|
|---|---|---|
|
||||||
|
| `--fps` | 1 | Frames per second to sample (5–10 recommended for tracking) |
|
||||||
|
| `--prob-threshold` | 0.5 | Minimum calibrated match probability |
|
||||||
|
| `--match-threshold` | — | Raw cosine similarity threshold (fallback) |
|
||||||
|
| `--extinction` | 5s | How long a track persists after last detection |
|
||||||
|
| `--track-alpha` | — | IoU vs. embedding weight in Hungarian assignment |
|
||||||
|
| `--track-min-iou` | — | Minimum IoU gate for spatial assignment |
|
||||||
|
| `--track-max-embed` | — | Maximum embedding distance gate |
|
||||||
|
| `--track-max-missing` | — | Frames a track survives without a detection |
|
||||||
|
| `--track-min-frames` | 3 | Observations before a track's mean embedding is used for matching |
|
||||||
|
|
||||||
|
### Gallery builder
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python3 scripts/make_gallery.py --tmdb-bearer <JWT> --movie-id <TMDB_ID> --output gallery.json
|
||||||
|
```
|
||||||
|
|
||||||
|
Fetches cast images from TMDB and embeds them via `embed_faces`.
|
||||||
|
|
||||||
|
## Output format
|
||||||
|
|
||||||
|
**Minimal** (default) — Jellyfin-ready:
|
||||||
|
```json
|
||||||
|
[
|
||||||
|
{ "actor": "Name", "start": 12.0, "end": 45.5 }
|
||||||
|
]
|
||||||
|
```
|
||||||
|
|
||||||
|
**Standard** — per-frame detail with bounding boxes, similarity scores, and track IDs.
|
||||||
|
|
||||||
|
## Pipeline topology
|
||||||
|
|
||||||
|
```
|
||||||
|
frame_source → face_detector → face_aligner → embedder
|
||||||
|
→ face_tracker → identity_matcher → scene_tracker → result_sink
|
||||||
|
```
|
||||||
|
|
||||||
|
Debug/preview branches fan out automatically from `identity_matcher`.
|
||||||
|
|
||||||
|
## Evaluation
|
||||||
|
|
||||||
|
Scripts in `eval/` and `scripts/movienet_*.py` support benchmarking against the MovieNet dataset.
|
||||||
+362
@@ -0,0 +1,362 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0000.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0119_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0001.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0125_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0002.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0128_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0003.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0225_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0004.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0306_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0005.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0308_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0006.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0311_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0007.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0315_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0008.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0317_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0001589_0009.jpg",
|
||||||
|
"imdb_id": "nm0001589",
|
||||||
|
"actor_name": "Michael Palin",
|
||||||
|
"source_frame": "tt0079470/shot_0319_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0000.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0072_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0001.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0074_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0002.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0078_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0003.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0079_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0004.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0080_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0005.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0638_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0006.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0101410/shot_0641_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0007.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0105236/shot_0017_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0008.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0105236/shot_0022_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0000114_0009.jpg",
|
||||||
|
"imdb_id": "nm0000114",
|
||||||
|
"actor_name": "Steve Buscemi",
|
||||||
|
"source_frame": "tt0105236/shot_0024_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0000.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0097_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0001.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0100_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0002.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0108_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0003.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0110_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0004.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0121_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0005.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0123_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0006.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0151_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0007.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0158_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0008.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0178_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0005042_0009.jpg",
|
||||||
|
"imdb_id": "nm0005042",
|
||||||
|
"actor_name": "Jason Isaacs",
|
||||||
|
"source_frame": "tt0119081/shot_0180_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0000.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0154_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0001.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0157_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0002.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0161_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0003.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0163_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0004.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0164_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0005.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0167_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0006.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0199_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0007.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0200_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0008.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0201_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm0175916_0009.jpg",
|
||||||
|
"imdb_id": "nm0175916",
|
||||||
|
"actor_name": "Paddy Considine",
|
||||||
|
"source_frame": "tt0440963/shot_0202_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0000.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0005_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0001.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0009_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0002.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0010_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0003.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0013_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0004.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0014_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0005.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0543_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0006.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0585_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0007.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0628_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0008.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0631_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm1385871_0009.jpg",
|
||||||
|
"imdb_id": "nm1385871",
|
||||||
|
"actor_name": "Olga Kurylenko",
|
||||||
|
"source_frame": "tt1483013/shot_0635_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0000.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0019_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0001.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0032_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0002.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0055_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0003.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0057_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0004.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0059_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0005.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0063_img_1.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0006.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0072_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0007.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0074_img_0.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0008.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0075_img_2.jpg"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"crop": "eval/probe/nm2057859_0009.jpg",
|
||||||
|
"imdb_id": "nm2057859",
|
||||||
|
"actor_name": "Andrea Riseborough",
|
||||||
|
"source_frame": "tt1483013/shot_0080_img_0.jpg"
|
||||||
|
}
|
||||||
|
]
|
||||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
|||||||
|
# MovieNet Validation Report
|
||||||
|
|
||||||
|
## Summary
|
||||||
|
|
||||||
|
| Model | Rank-1 | Det.Fail | Mean-sim | Probes | Size |
|
||||||
|
| ----- | ------ | -------- | -------- | ------ | ----- |
|
||||||
|
| R50 | 85.2% | 10.0% | 0.470 | 60 | 167MB |
|
||||||
|
| R18 | 72.2% | 10.0% | 0.453 | 60 | 46MB |
|
||||||
|
| MBF | 83.3% | 10.0% | 0.383 | 60 | 13MB |
|
||||||
|
|
||||||
|
## Per-Actor Recall
|
||||||
|
|
||||||
|
| Actor | R50 | R18 | MBF |
|
||||||
|
| ------------------ | ------------ | ------------ | ------------ |
|
||||||
|
| Steve Buscemi | 100% (9/9) | 100% (9/9) | 100% (9/9) |
|
||||||
|
| Michael Palin | 88% (7/8) | 50% (4/8) | 62% (5/8) |
|
||||||
|
| Jason Isaacs | 100% (10/10) | 100% (10/10) | 100% (10/10) |
|
||||||
|
| Paddy Considine | 50% (5/10) | 50% (5/10) | 60% (6/10) |
|
||||||
|
| Olga Kurylenko | 89% (8/9) | 67% (6/9) | 100% (9/9) |
|
||||||
|
| Andrea Riseborough | 88% (7/8) | 62% (5/8) | 75% (6/8) |
|
||||||
+1
Submodule external/KPN added at 79916f1da1
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Executable
+51
@@ -0,0 +1,51 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
# Pre-build TensorRT engines for ArcFace and SCRFD with the same shape profiles
|
||||||
|
# the runtime nodes use. First-run ORT engine builds take 30–90 s per model and
|
||||||
|
# block the pipeline; this script does it offline so cold starts are instant.
|
||||||
|
#
|
||||||
|
# Profiles must match src/arcface_embedder.hpp and src/scrfd_decoder.hpp:
|
||||||
|
# ArcFace : min=1x3x112x112 opt=Nx3x112x112 max=Nx3x112x112 (N = embed batch)
|
||||||
|
# SCRFD : 1x3x640x640 (fixed; we letterbox to this)
|
||||||
|
#
|
||||||
|
# These trtexec-built engines are *not* picked up by the ORT TRT EP cache —
|
||||||
|
# ORT uses its own engine format. The point of this script is:
|
||||||
|
# (a) sanity-check that the ONNX models build under TRT at all;
|
||||||
|
# (b) measure pure inference latency without ORT overhead.
|
||||||
|
# Run scene_analyze normally and ORT will populate ./trt_cache itself.
|
||||||
|
set -euo pipefail
|
||||||
|
|
||||||
|
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||||
|
MODELS="$ROOT/models"
|
||||||
|
OUT="$ROOT/trt_cache"
|
||||||
|
mkdir -p "$OUT"
|
||||||
|
|
||||||
|
EMBED_BATCH="${EMBED_BATCH:-4}"
|
||||||
|
ARCFACE_MODEL="${ARCFACE_MODEL:-$MODELS/arcface_w600k_r50.onnx}"
|
||||||
|
SCRFD_MODEL="${SCRFD_MODEL:-$MODELS/scrfd_500m_bnkps.onnx}"
|
||||||
|
|
||||||
|
run() { echo "+ $*"; "$@"; }
|
||||||
|
|
||||||
|
echo "== ArcFace =="
|
||||||
|
run trtexec \
|
||||||
|
--onnx="$ARCFACE_MODEL" \
|
||||||
|
--fp16 \
|
||||||
|
--minShapes=input.1:1x3x112x112 \
|
||||||
|
--optShapes=input.1:${EMBED_BATCH}x3x112x112 \
|
||||||
|
--maxShapes=input.1:${EMBED_BATCH}x3x112x112 \
|
||||||
|
--saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \
|
||||||
|
--useCudaGraph
|
||||||
|
|
||||||
|
echo
|
||||||
|
echo "== SCRFD =="
|
||||||
|
run trtexec \
|
||||||
|
--onnx="$SCRFD_MODEL" \
|
||||||
|
--fp16 \
|
||||||
|
--minShapes=input.1:1x3x640x640 \
|
||||||
|
--optShapes=input.1:1x3x640x640 \
|
||||||
|
--maxShapes=input.1:1x3x640x640 \
|
||||||
|
--saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \
|
||||||
|
--useCudaGraph
|
||||||
|
|
||||||
|
echo
|
||||||
|
echo "Engines saved under: $OUT"
|
||||||
|
echo "Look for 'mean: ... ms' in each section for per-call latency."
|
||||||
Executable
+71
@@ -0,0 +1,71 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
# Download ONNX models required by scene_analyze and build_gallery.
|
||||||
|
# Run from the project root: bash scripts/download_models.sh
|
||||||
|
|
||||||
|
set -euo pipefail
|
||||||
|
|
||||||
|
MODELS_DIR="${1:-models}"
|
||||||
|
mkdir -p "$MODELS_DIR"
|
||||||
|
|
||||||
|
# ── YuNet face detection ──────────────────────────────────────────────────────
|
||||||
|
YUNET_URL="https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
|
||||||
|
YUNET_FILE="$MODELS_DIR/face_detection_yunet_2023mar.onnx"
|
||||||
|
if [ ! -f "$YUNET_FILE" ]; then
|
||||||
|
echo "Downloading YuNet…"
|
||||||
|
curl -L "$YUNET_URL" -o "$YUNET_FILE"
|
||||||
|
else
|
||||||
|
echo "YuNet already present: $YUNET_FILE"
|
||||||
|
fi
|
||||||
|
|
||||||
|
# ── ArcFace face recognition (buffalo_l / w600k_r50) ─────────────────────────
|
||||||
|
# This model is part of InsightFace's buffalo_l pack.
|
||||||
|
# We download and unpack only the recognition model.
|
||||||
|
ARCFACE_FILE="$MODELS_DIR/arcface_w600k_r50.onnx"
|
||||||
|
if [ ! -f "$ARCFACE_FILE" ]; then
|
||||||
|
echo "Downloading ArcFace (buffalo_l)…"
|
||||||
|
TMP_ZIP=$(mktemp /tmp/buffalo_l.XXXXXX.zip)
|
||||||
|
curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip" \
|
||||||
|
-o "$TMP_ZIP"
|
||||||
|
# The zip contains: 1k3d68.onnx 2d106det.onnx det_10g.onnx genderage.onnx w600k_r50.onnx
|
||||||
|
unzip -jo "$TMP_ZIP" "w600k_r50.onnx" -d "$MODELS_DIR"
|
||||||
|
mv "$MODELS_DIR/w600k_r50.onnx" "$ARCFACE_FILE"
|
||||||
|
rm "$TMP_ZIP"
|
||||||
|
else
|
||||||
|
echo "ArcFace already present: $ARCFACE_FILE"
|
||||||
|
fi
|
||||||
|
|
||||||
|
# ── ArcFace face recognition (buffalo_s / w600k_mbf — MobileFaceNet) ─────────
|
||||||
|
# Lighter backbone (13 MB vs 174 MB for R50) — same 512-dim output, faster inference.
|
||||||
|
ARCFACE_MBF_FILE="$MODELS_DIR/arcface_w600k_mbf.onnx"
|
||||||
|
if [ ! -f "$ARCFACE_MBF_FILE" ]; then
|
||||||
|
echo "Downloading ArcFace MobileFaceNet (buffalo_s)…"
|
||||||
|
TMP_ZIP=$(mktemp /tmp/buffalo_s.XXXXXX.zip)
|
||||||
|
curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_s.zip" \
|
||||||
|
-o "$TMP_ZIP"
|
||||||
|
unzip -jo "$TMP_ZIP" "w600k_mbf.onnx" -d "$MODELS_DIR"
|
||||||
|
mv "$MODELS_DIR/w600k_mbf.onnx" "$ARCFACE_MBF_FILE"
|
||||||
|
rm "$TMP_ZIP"
|
||||||
|
else
|
||||||
|
echo "ArcFace MBF already present: $ARCFACE_MBF_FILE"
|
||||||
|
fi
|
||||||
|
|
||||||
|
# ── SCRFD-500MF face detection (InsightFace buffalo_sc) ───────────────────────
|
||||||
|
# buffalo_sc.zip contains det_500m.onnx (SCRFD-500MF with 5 keypoints).
|
||||||
|
# If the unzip fails (file not found in archive), download manually from:
|
||||||
|
# https://huggingface.co/deepinsight/insightface/resolve/main/models/buffalo_sc/det_500m.onnx
|
||||||
|
SCRFD_FILE="$MODELS_DIR/scrfd_500m_bnkps.onnx"
|
||||||
|
if [ ! -f "$SCRFD_FILE" ]; then
|
||||||
|
echo "Downloading SCRFD-500MF (buffalo_sc)…"
|
||||||
|
TMP_ZIP=$(mktemp /tmp/buffalo_sc.XXXXXX.zip)
|
||||||
|
curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_sc.zip" \
|
||||||
|
-o "$TMP_ZIP"
|
||||||
|
unzip -jo "$TMP_ZIP" "det_500m.onnx" -d "$MODELS_DIR"
|
||||||
|
mv "$MODELS_DIR/det_500m.onnx" "$SCRFD_FILE"
|
||||||
|
rm "$TMP_ZIP"
|
||||||
|
else
|
||||||
|
echo "SCRFD-500MF already present: $SCRFD_FILE"
|
||||||
|
fi
|
||||||
|
|
||||||
|
echo ""
|
||||||
|
echo "Models ready in $MODELS_DIR/:"
|
||||||
|
ls -lh "$MODELS_DIR"
|
||||||
Executable
+291
@@ -0,0 +1,291 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""make_gallery.py — fetch actor images for a movie and build gallery.json.
|
||||||
|
|
||||||
|
Fetches the cast from TMDB, downloads actor profile images, runs the C++
|
||||||
|
embed_faces binary (SCRFD + ArcFace, same models as scene_analyze) to produce
|
||||||
|
embeddings, then writes gallery.json.
|
||||||
|
|
||||||
|
Requirements:
|
||||||
|
pip install requests Pillow
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
# By IMDB movie ID (most natural — resolves to TMDB automatically):
|
||||||
|
python scripts/make_gallery.py \\
|
||||||
|
--tmdb-key YOUR_KEY \\
|
||||||
|
--imdb-id tt0137523 \\
|
||||||
|
--output gallery.json
|
||||||
|
|
||||||
|
# Or directly with a TMDB movie ID:
|
||||||
|
python scripts/make_gallery.py \\
|
||||||
|
--tmdb-key YOUR_KEY \\
|
||||||
|
--movie-id 550 \\
|
||||||
|
--output gallery.json
|
||||||
|
|
||||||
|
# Additional options:
|
||||||
|
# --embed-bin build/embed_faces path to embed_faces binary
|
||||||
|
# --models-dir models/ directory with ONNX models
|
||||||
|
# --max-actors 20 how many cast members to include
|
||||||
|
# --images-per-actor 3 profile images to download per actor
|
||||||
|
# --image-dir /tmp/gallery_imgs where to cache downloaded images
|
||||||
|
|
||||||
|
Get a free TMDB API key at: https://www.themoviedb.org/settings/api
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import io
|
||||||
|
|
||||||
|
import requests
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
TMDB_BASE = "https://api.themoviedb.org/3"
|
||||||
|
TMDB_IMG = "https://image.tmdb.org/t/p/original"
|
||||||
|
|
||||||
|
# ── TMDB helpers ──────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def tmdb_get(path: str, token: str, **params) -> dict:
|
||||||
|
url = TMDB_BASE + path
|
||||||
|
if token.startswith("eyJ"):
|
||||||
|
headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"}
|
||||||
|
r = requests.get(url, params=params, headers=headers, timeout=10)
|
||||||
|
else:
|
||||||
|
params["api_key"] = token
|
||||||
|
r = requests.get(url, params=params, headers={"Accept": "application/json"}, timeout=10)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
def tmdb_id_from_imdb(imdb_id: str, key: str) -> int:
|
||||||
|
data = tmdb_get(f"/find/{imdb_id}", key, external_source="imdb_id")
|
||||||
|
results = data.get("movie_results", [])
|
||||||
|
if not results:
|
||||||
|
raise ValueError(f"No TMDB movie found for IMDB ID {imdb_id}")
|
||||||
|
return results[0]["id"]
|
||||||
|
|
||||||
|
|
||||||
|
def fetch_cast(movie_id: int, key: str, max_actors: int) -> list[dict]:
|
||||||
|
"""Return list of {id, name, imdb_id, profile_images: [...url...]}."""
|
||||||
|
credits = tmdb_get(f"/movie/{movie_id}/credits", key)
|
||||||
|
cast = credits.get("cast", [])[:max_actors]
|
||||||
|
|
||||||
|
actors = []
|
||||||
|
for member in cast:
|
||||||
|
person_id = member["id"]
|
||||||
|
|
||||||
|
# Get IMDB ID for this person
|
||||||
|
ext = tmdb_get(f"/person/{person_id}/external_ids", key)
|
||||||
|
imdb_id = ext.get("imdb_id") or f"tmdb_{person_id}"
|
||||||
|
|
||||||
|
# Get profile images (sorted by vote_average desc by TMDB)
|
||||||
|
images_data = tmdb_get(f"/person/{person_id}/images", key)
|
||||||
|
profiles = images_data.get("profiles", [])
|
||||||
|
image_urls = [TMDB_IMG + p["file_path"] for p in profiles if p.get("file_path")]
|
||||||
|
|
||||||
|
if not image_urls:
|
||||||
|
print(f" [warn] no images for {member['name']}, skipping", file=sys.stderr)
|
||||||
|
continue
|
||||||
|
|
||||||
|
actors.append({
|
||||||
|
"id": person_id,
|
||||||
|
"name": member["name"],
|
||||||
|
"imdb_id": imdb_id,
|
||||||
|
"profile_images": image_urls,
|
||||||
|
})
|
||||||
|
time.sleep(0.05) # be polite to TMDB
|
||||||
|
|
||||||
|
return actors
|
||||||
|
|
||||||
|
|
||||||
|
# ── Image download ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def download_images(actor: dict, dest_dir: Path, n: int) -> list[Path]:
|
||||||
|
"""Download up to n profile images for an actor into dest_dir."""
|
||||||
|
dest_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
paths = []
|
||||||
|
for i, url in enumerate(actor["profile_images"][:n]):
|
||||||
|
out = dest_dir / f"{i:02d}.jpg"
|
||||||
|
if out.exists() and out.stat().st_size > 1024:
|
||||||
|
paths.append(out)
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
r = requests.get(url, timeout=15)
|
||||||
|
r.raise_for_status()
|
||||||
|
Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG")
|
||||||
|
paths.append(out)
|
||||||
|
except Exception as e:
|
||||||
|
print(f" [warn] download failed: {url}: {e}", file=sys.stderr)
|
||||||
|
return paths
|
||||||
|
|
||||||
|
|
||||||
|
# ── Embedding via embed_faces binary ─────────────────────────────────────────
|
||||||
|
|
||||||
|
def embed_images(image_paths: list[Path], embed_bin: str,
|
||||||
|
detector: str, arcface: str) -> list[dict | None]:
|
||||||
|
"""
|
||||||
|
Call the C++ embed_faces binary on a list of images.
|
||||||
|
Returns a list of result dicts (or None if no face / error) per image.
|
||||||
|
"""
|
||||||
|
if not image_paths:
|
||||||
|
return []
|
||||||
|
|
||||||
|
cmd = [
|
||||||
|
embed_bin,
|
||||||
|
"--detector", detector,
|
||||||
|
"--arcface", arcface,
|
||||||
|
] + [str(p) for p in image_paths]
|
||||||
|
|
||||||
|
try:
|
||||||
|
proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
||||||
|
except subprocess.CalledProcessError as e:
|
||||||
|
print(f"[error] embed_faces failed:\n{e.stderr}", file=sys.stderr)
|
||||||
|
return [None] * len(image_paths)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = json.loads(proc.stdout)
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
print(f"[error] embed_faces output is not valid JSON: {e}", file=sys.stderr)
|
||||||
|
return [None] * len(image_paths)
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
# ── Gallery assembly ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def build_gallery(movie_id: int, key: str, embed_bin: str,
|
||||||
|
detector: str, arcface: str,
|
||||||
|
max_actors: int, images_per_actor: int,
|
||||||
|
image_root: Path) -> dict:
|
||||||
|
"""Fetch cast, download images, embed, return gallery dict."""
|
||||||
|
print(f"Fetching cast for TMDB movie {movie_id}…", file=sys.stderr)
|
||||||
|
actors = fetch_cast(movie_id, key, max_actors)
|
||||||
|
print(f"Found {len(actors)} actors with images", file=sys.stderr)
|
||||||
|
|
||||||
|
gallery_actors = []
|
||||||
|
|
||||||
|
for actor in actors:
|
||||||
|
safe_name = actor["name"].replace(" ", "_")
|
||||||
|
actor_dir = image_root / f"{actor['imdb_id']}_{safe_name}"
|
||||||
|
|
||||||
|
print(f"\n{actor['name']} ({actor['imdb_id']})", file=sys.stderr)
|
||||||
|
image_paths = download_images(actor, actor_dir, images_per_actor)
|
||||||
|
if not image_paths:
|
||||||
|
print(" no images downloaded, skipping", file=sys.stderr)
|
||||||
|
continue
|
||||||
|
|
||||||
|
print(f" embedding {len(image_paths)} image(s)…", file=sys.stderr)
|
||||||
|
results = embed_images(image_paths, embed_bin, detector, arcface)
|
||||||
|
|
||||||
|
embeddings = []
|
||||||
|
source_images = []
|
||||||
|
for path, res in zip(image_paths, results):
|
||||||
|
if res is None or res.get("embedding") is None:
|
||||||
|
reason = res.get("error", "unknown") if res else "binary error"
|
||||||
|
print(f" [skip] {path.name}: {reason}", file=sys.stderr)
|
||||||
|
continue
|
||||||
|
embeddings.append(res["embedding"])
|
||||||
|
source_images.append(path.name)
|
||||||
|
print(f" [ok] {path.name} conf={res.get('confidence', 0):.2f}",
|
||||||
|
file=sys.stderr)
|
||||||
|
|
||||||
|
if not embeddings:
|
||||||
|
print(" no valid embeddings, skipping actor", file=sys.stderr)
|
||||||
|
continue
|
||||||
|
|
||||||
|
gallery_actors.append({
|
||||||
|
"imdb_id": actor["imdb_id"],
|
||||||
|
"name": actor["name"],
|
||||||
|
"source_images": source_images,
|
||||||
|
"embeddings": embeddings,
|
||||||
|
})
|
||||||
|
print(f" → {len(embeddings)} embedding(s) stored", file=sys.stderr)
|
||||||
|
|
||||||
|
return {"actors": gallery_actors}
|
||||||
|
|
||||||
|
|
||||||
|
# ── Entry point ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Fetch TMDB cast images and build gallery.json via embed_faces")
|
||||||
|
parser.add_argument("--tmdb-key", required=True,
|
||||||
|
help="TMDB Bearer token (API Read Access Token from themoviedb.org/settings/api)")
|
||||||
|
group = parser.add_mutually_exclusive_group(required=True)
|
||||||
|
group.add_argument("--imdb-id",
|
||||||
|
help="IMDB movie ID, e.g. tt0137523 — looked up via TMDB automatically")
|
||||||
|
group.add_argument("--movie-id", type=int,
|
||||||
|
help="TMDB movie ID (alternative to --imdb-id)")
|
||||||
|
parser.add_argument("--output", required=True, help="Output gallery.json path")
|
||||||
|
parser.add_argument("--embed-bin", default="build/embed_faces",
|
||||||
|
help="Path to embed_faces binary (default: build/embed_faces)")
|
||||||
|
parser.add_argument("--models-dir", default="models",
|
||||||
|
help="Directory containing ONNX models (default: models/)")
|
||||||
|
parser.add_argument("--arcface", default=None,
|
||||||
|
help="Path to ArcFace ONNX model (overrides --models-dir selection)")
|
||||||
|
parser.add_argument("--max-actors", type=int, default=20,
|
||||||
|
help="Maximum number of cast members to include (default: 20)")
|
||||||
|
parser.add_argument("--images-per-actor",type=int, default=3,
|
||||||
|
help="Profile images to download per actor (default: 3)")
|
||||||
|
parser.add_argument("--image-dir", default=None,
|
||||||
|
help="Where to store downloaded images (default: <output_dir>/images)")
|
||||||
|
parser.add_argument("--keep-images", action="store_true",
|
||||||
|
help="Do not delete downloaded images after embedding")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# Resolve paths
|
||||||
|
embed_bin = str(Path(args.embed_bin).resolve())
|
||||||
|
models_dir = Path(args.models_dir)
|
||||||
|
detector = str(models_dir / "scrfd_500m_bnkps.onnx")
|
||||||
|
arcface = args.arcface if args.arcface else str(models_dir / "arcface_w600k_r50.onnx")
|
||||||
|
output = Path(args.output)
|
||||||
|
image_root = Path(args.image_dir) if args.image_dir else output.parent / "images"
|
||||||
|
|
||||||
|
# Validate
|
||||||
|
if not Path(embed_bin).is_file():
|
||||||
|
sys.exit(f"embed_faces binary not found: {embed_bin}\n"
|
||||||
|
f"Build it first: cmake --build build --target embed_faces")
|
||||||
|
for model, name in [(detector, "SCRFD"), (arcface, "ArcFace")]:
|
||||||
|
if not Path(model).is_file():
|
||||||
|
sys.exit(f"{name} model not found: {model}\n"
|
||||||
|
f"Run: bash scripts/download_models.sh")
|
||||||
|
|
||||||
|
# Resolve movie ID
|
||||||
|
movie_id = args.movie_id
|
||||||
|
if movie_id is None:
|
||||||
|
print(f"Resolving IMDB ID {args.imdb_id} → TMDB…", file=sys.stderr)
|
||||||
|
movie_id = tmdb_id_from_imdb(args.imdb_id, args.tmdb_key)
|
||||||
|
print(f"TMDB movie ID: {movie_id}", file=sys.stderr)
|
||||||
|
|
||||||
|
# Build gallery
|
||||||
|
gallery = build_gallery(
|
||||||
|
movie_id = movie_id,
|
||||||
|
key = args.tmdb_key,
|
||||||
|
embed_bin = embed_bin,
|
||||||
|
detector = detector,
|
||||||
|
arcface = arcface,
|
||||||
|
max_actors = args.max_actors,
|
||||||
|
images_per_actor = args.images_per_actor,
|
||||||
|
image_root = image_root,
|
||||||
|
)
|
||||||
|
|
||||||
|
n_actors = len(gallery["actors"])
|
||||||
|
n_embeddings = sum(len(a["embeddings"]) for a in gallery["actors"])
|
||||||
|
print(f"\nGallery: {n_actors} actors, {n_embeddings} total embeddings",
|
||||||
|
file=sys.stderr)
|
||||||
|
|
||||||
|
if n_actors == 0:
|
||||||
|
sys.exit("No actors could be processed — check models and images.")
|
||||||
|
|
||||||
|
output.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
output.write_text(json.dumps(gallery, indent=2) + "\n")
|
||||||
|
print(f"Saved: {output}", file=sys.stderr)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,144 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
movienet_eval.py — embed probe crops and match against a gallery.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python scripts/movienet_eval.py \
|
||||||
|
--gallery gallery_r50.json \
|
||||||
|
--arcface models/arcface_w600k_r50.onnx \
|
||||||
|
--gt eval/gt.json \
|
||||||
|
--output eval/predictions_r50.json \
|
||||||
|
[--yunet models/face_detection_yunet_2023mar.onnx] \
|
||||||
|
[--embed-bin build/embed_faces]
|
||||||
|
|
||||||
|
Input (--gt): list of {"crop": <path>, "imdb_id": <str>, "actor_name": <str>}
|
||||||
|
Output: list of {"crop", "gt", "pred", "similarity", "detection_failed", "all_scores"}
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def load_gallery(path: str) -> dict[str, dict]:
|
||||||
|
"""Return {imdb_id: {"name": str, "embeddings": [[float]]}}."""
|
||||||
|
with open(path) as f:
|
||||||
|
data = json.load(f)
|
||||||
|
return {a["imdb_id"]: {"name": a["name"], "embeddings": a["embeddings"]}
|
||||||
|
for a in data["actors"]}
|
||||||
|
|
||||||
|
|
||||||
|
def dot(a: list[float], b: list[float]) -> float:
|
||||||
|
return sum(x * y for x, y in zip(a, b))
|
||||||
|
|
||||||
|
|
||||||
|
def embed_images(paths: list[Path], embed_bin: str, yunet: str, arcface: str) -> list[dict | None]:
|
||||||
|
if not paths:
|
||||||
|
return []
|
||||||
|
cmd = [embed_bin, "--yunet", yunet, "--arcface", arcface] + [str(p) for p in paths]
|
||||||
|
try:
|
||||||
|
proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
||||||
|
except subprocess.CalledProcessError as e:
|
||||||
|
print(f"[error] embed_faces failed:\n{e.stderr}", file=sys.stderr)
|
||||||
|
return [None] * len(paths)
|
||||||
|
try:
|
||||||
|
return json.loads(proc.stdout)
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
print(f"[error] embed_faces JSON parse error: {e}", file=sys.stderr)
|
||||||
|
return [None] * len(paths)
|
||||||
|
|
||||||
|
|
||||||
|
def match(embedding: list[float], gallery: dict[str, dict]) -> tuple[str, float, dict[str, float]]:
|
||||||
|
"""Return (best_imdb_id, best_similarity, {imdb_id: similarity})."""
|
||||||
|
scores: dict[str, float] = {}
|
||||||
|
for imdb_id, actor in gallery.items():
|
||||||
|
# max similarity across all reference embeddings for this actor
|
||||||
|
scores[imdb_id] = max(dot(embedding, ref) for ref in actor["embeddings"])
|
||||||
|
best_id = max(scores, key=lambda k: scores[k])
|
||||||
|
return best_id, scores[best_id], scores
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p = argparse.ArgumentParser()
|
||||||
|
p.add_argument("--gallery", required=True)
|
||||||
|
p.add_argument("--arcface", required=True)
|
||||||
|
p.add_argument("--gt", required=True)
|
||||||
|
p.add_argument("--output", required=True)
|
||||||
|
p.add_argument("--yunet", default="models/face_detection_yunet_2023mar.onnx")
|
||||||
|
p.add_argument("--embed-bin", default="build/embed_faces")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
gallery = load_gallery(args.gallery)
|
||||||
|
print(f"[eval] gallery: {len(gallery)} actors", file=sys.stderr)
|
||||||
|
|
||||||
|
with open(args.gt) as f:
|
||||||
|
gt_entries = json.load(f)
|
||||||
|
print(f"[eval] probe crops: {len(gt_entries)}", file=sys.stderr)
|
||||||
|
|
||||||
|
# Batch all crops in one embed_faces call to amortise startup cost
|
||||||
|
crop_paths = [Path(e["crop"]) for e in gt_entries]
|
||||||
|
missing = [p for p in crop_paths if not p.exists()]
|
||||||
|
if missing:
|
||||||
|
print(f"[warn] {len(missing)} crop(s) not found on disk, skipping", file=sys.stderr)
|
||||||
|
|
||||||
|
results_raw = embed_images(
|
||||||
|
[p for p in crop_paths if p.exists()],
|
||||||
|
args.embed_bin, args.yunet, args.arcface
|
||||||
|
)
|
||||||
|
|
||||||
|
# Re-index results back to original list (missing files get None)
|
||||||
|
raw_iter = iter(results_raw)
|
||||||
|
embed_results: list[dict | None] = []
|
||||||
|
for p in crop_paths:
|
||||||
|
embed_results.append(next(raw_iter) if p.exists() else None)
|
||||||
|
|
||||||
|
predictions = []
|
||||||
|
n_det_fail = 0
|
||||||
|
n_correct = 0
|
||||||
|
|
||||||
|
for entry, result in zip(gt_entries, embed_results):
|
||||||
|
detection_failed = result is None or result.get("embedding") is None
|
||||||
|
if detection_failed:
|
||||||
|
n_det_fail += 1
|
||||||
|
predictions.append({
|
||||||
|
"crop": entry["crop"],
|
||||||
|
"gt": entry["imdb_id"],
|
||||||
|
"pred": None,
|
||||||
|
"similarity": None,
|
||||||
|
"detection_failed": True,
|
||||||
|
"all_scores": {},
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
|
||||||
|
pred_id, sim, all_scores = match(result["embedding"], gallery)
|
||||||
|
correct = pred_id == entry["imdb_id"]
|
||||||
|
if correct:
|
||||||
|
n_correct += 1
|
||||||
|
|
||||||
|
predictions.append({
|
||||||
|
"crop": entry["crop"],
|
||||||
|
"gt": entry["imdb_id"],
|
||||||
|
"pred": pred_id,
|
||||||
|
"similarity": sim,
|
||||||
|
"detection_failed": False,
|
||||||
|
"all_scores": all_scores,
|
||||||
|
})
|
||||||
|
|
||||||
|
n_total = len(gt_entries)
|
||||||
|
n_evaluated = n_total - n_det_fail
|
||||||
|
rank1 = n_correct / n_evaluated * 100 if n_evaluated else 0
|
||||||
|
print(f"[eval] detection failures: {n_det_fail}/{n_total}", file=sys.stderr)
|
||||||
|
print(f"[eval] rank-1 accuracy: {rank1:.1f}% ({n_correct}/{n_evaluated})", file=sys.stderr)
|
||||||
|
|
||||||
|
out_path = Path(args.output)
|
||||||
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with open(out_path, "w") as f:
|
||||||
|
json.dump(predictions, f, indent=2)
|
||||||
|
print(f"[eval] written → {out_path}", file=sys.stderr)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,199 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
movienet_prep.py — extract probe crops from MovieNet-PS for actors in our gallery.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python scripts/movienet_prep.py \
|
||||||
|
--movienet <movienet_root> \
|
||||||
|
--gallery gallery.json \
|
||||||
|
--output eval/ \
|
||||||
|
[--split Train_app10] \
|
||||||
|
[--margin 0.2] \
|
||||||
|
[--max-per-actor 50]
|
||||||
|
|
||||||
|
MovieNet-PS format (annotation.zip + Image.zip):
|
||||||
|
annotation/test/train_test/Train_app<N>.mat — N annotations per actor
|
||||||
|
Train[i] = [imdb_id (nm...), count, [[img_path, bbox[x,y,w,h], label], ...]]
|
||||||
|
Image/<movie_tt_id>/shot_XXXX_img_Y.jpg — source frames
|
||||||
|
|
||||||
|
Output:
|
||||||
|
eval/probe/ face crops (jpg)
|
||||||
|
eval/gt.json [{crop, imdb_id, actor_name, source_frame}]
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import zipfile
|
||||||
|
from io import BytesIO
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
try:
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import scipy.io as sio
|
||||||
|
except ImportError as e:
|
||||||
|
print(f"[error] missing dependency: {e}", file=sys.stderr)
|
||||||
|
print("Install: pip install opencv-python scipy numpy", file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
# ── MovieNet-PS loader ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def load_movienet_annotations(movienet_root: Path, split: str) -> list[dict]:
|
||||||
|
"""
|
||||||
|
Parse a MovieNet-PS Train_app<N>.mat split.
|
||||||
|
Returns flat list of {"imdb_id", "img_path", "bbox": [x,y,w,h]}.
|
||||||
|
img_path is relative to Image/ inside Image.zip, e.g. tt0047396/shot_0004_img_1.jpg
|
||||||
|
"""
|
||||||
|
mat_path = movienet_root / "annotation" / "test" / "train_test" / f"{split}.mat"
|
||||||
|
if not mat_path.exists():
|
||||||
|
# try extracting from annotation.zip
|
||||||
|
zip_path = movienet_root / "annotation.zip"
|
||||||
|
if not zip_path.exists():
|
||||||
|
raise FileNotFoundError(f"annotation.zip not found in {movienet_root}")
|
||||||
|
inner = f"annotation/test/train_test/{split}.mat"
|
||||||
|
print(f"[prep] extracting {inner} from annotation.zip…", file=sys.stderr)
|
||||||
|
with zipfile.ZipFile(zip_path) as z:
|
||||||
|
z.extract(inner, movienet_root)
|
||||||
|
mat_path = movienet_root / inner
|
||||||
|
|
||||||
|
data = sio.loadmat(str(mat_path))["Train"]
|
||||||
|
annotations = []
|
||||||
|
for row in data:
|
||||||
|
imdb_id = str(row[0].flat[0]) # e.g. "nm0000023"
|
||||||
|
entries = row[2] # array of [path, bbox, label]
|
||||||
|
for entry in entries:
|
||||||
|
img_path = str(entry[0].flat[0]) # e.g. "tt0032138/shot_0003_img_1.jpg"
|
||||||
|
bbox = [float(v) for v in entry[1].flat] # [x, y, w, h]
|
||||||
|
annotations.append({"imdb_id": imdb_id, "img_path": img_path, "bbox": bbox})
|
||||||
|
return annotations
|
||||||
|
|
||||||
|
|
||||||
|
# ── Gallery loader ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def load_gallery_ids(gallery_path: str) -> dict[str, str]:
|
||||||
|
"""Return {imdb_id: actor_name} for all actors in the gallery."""
|
||||||
|
with open(gallery_path) as f:
|
||||||
|
data = json.load(f)
|
||||||
|
return {a["imdb_id"]: a["name"] for a in data["actors"]}
|
||||||
|
|
||||||
|
|
||||||
|
# ── Crop + save ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def crop_face(img: "np.ndarray", bbox: list[float], margin: float) -> "np.ndarray | None":
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
x, y, bw, bh = bbox
|
||||||
|
# expand by margin
|
||||||
|
pad_x = bw * margin
|
||||||
|
pad_y = bh * margin
|
||||||
|
x1 = max(0, int(x - pad_x))
|
||||||
|
y1 = max(0, int(y - pad_y))
|
||||||
|
x2 = min(w, int(x + bw + pad_x))
|
||||||
|
y2 = min(h, int(y + bh + pad_y))
|
||||||
|
crop = img[y1:y2, x1:x2]
|
||||||
|
return crop if crop.size > 0 else None
|
||||||
|
|
||||||
|
|
||||||
|
# ── Main ──────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p = argparse.ArgumentParser()
|
||||||
|
p.add_argument("--movienet", required=True, help="MovieNet-PS root directory")
|
||||||
|
p.add_argument("--gallery", required=True, help="gallery.json (for actor list)")
|
||||||
|
p.add_argument("--output", default="eval", help="output directory")
|
||||||
|
p.add_argument("--split", default="Train_app10",
|
||||||
|
help="annotation split to use (default: Train_app10)")
|
||||||
|
p.add_argument("--margin", type=float, default=0.2,
|
||||||
|
help="bbox expansion factor (default 0.2 = 20%%)")
|
||||||
|
p.add_argument("--max-per-actor", type=int, default=50,
|
||||||
|
help="cap probe crops per actor (default 50)")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
movienet_root = Path(args.movienet)
|
||||||
|
out_dir = Path(args.output)
|
||||||
|
probe_dir = out_dir / "probe"
|
||||||
|
probe_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
gallery_ids = load_gallery_ids(args.gallery)
|
||||||
|
print(f"[prep] gallery actors: {len(gallery_ids)}", file=sys.stderr)
|
||||||
|
|
||||||
|
annotations = load_movienet_annotations(movienet_root, args.split)
|
||||||
|
print(f"[prep] total annotations in split: {len(annotations)}", file=sys.stderr)
|
||||||
|
|
||||||
|
matched = [a for a in annotations if a["imdb_id"] in gallery_ids]
|
||||||
|
print(f"[prep] annotations matching gallery: {len(matched)}", file=sys.stderr)
|
||||||
|
|
||||||
|
if not matched:
|
||||||
|
print("[error] no overlap between MovieNet and gallery — check IMDb ID format", file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
# Build set of image paths we actually need, then extract from Image.zip in one pass
|
||||||
|
needed_paths = {a["img_path"] for a in matched}
|
||||||
|
image_zip = movienet_root / "Image.zip"
|
||||||
|
frame_cache: dict[str, np.ndarray] = {}
|
||||||
|
|
||||||
|
print(f"[prep] extracting {len(needed_paths)} frames from Image.zip…", file=sys.stderr)
|
||||||
|
with zipfile.ZipFile(image_zip) as zf:
|
||||||
|
for img_path in needed_paths:
|
||||||
|
zip_entry = f"Image/{img_path}"
|
||||||
|
try:
|
||||||
|
data = zf.read(zip_entry)
|
||||||
|
arr = np.frombuffer(data, dtype=np.uint8)
|
||||||
|
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||||
|
if img is not None:
|
||||||
|
frame_cache[img_path] = img
|
||||||
|
except KeyError:
|
||||||
|
pass # file missing from zip, skip silently
|
||||||
|
|
||||||
|
print(f"[prep] frames loaded: {len(frame_cache)}/{len(needed_paths)}", file=sys.stderr)
|
||||||
|
|
||||||
|
per_actor_count: dict[str, int] = {}
|
||||||
|
gt_entries = []
|
||||||
|
n_failed = 0
|
||||||
|
|
||||||
|
for ann in matched:
|
||||||
|
imdb_id = ann["imdb_id"]
|
||||||
|
img_path = ann["img_path"]
|
||||||
|
count = per_actor_count.get(imdb_id, 0)
|
||||||
|
if count >= args.max_per_actor:
|
||||||
|
continue
|
||||||
|
|
||||||
|
img = frame_cache.get(img_path)
|
||||||
|
if img is None:
|
||||||
|
n_failed += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
crop = crop_face(img, ann["bbox"], args.margin)
|
||||||
|
if crop is None:
|
||||||
|
n_failed += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
crop_name = f"{imdb_id}_{count:04d}.jpg"
|
||||||
|
crop_path = probe_dir / crop_name
|
||||||
|
cv2.imwrite(str(crop_path), crop)
|
||||||
|
|
||||||
|
per_actor_count[imdb_id] = count + 1
|
||||||
|
gt_entries.append({
|
||||||
|
"crop": str(crop_path),
|
||||||
|
"imdb_id": imdb_id,
|
||||||
|
"actor_name": gallery_ids[imdb_id],
|
||||||
|
"source_frame": img_path,
|
||||||
|
})
|
||||||
|
|
||||||
|
gt_path = out_dir / "gt.json"
|
||||||
|
with open(gt_path, "w") as f:
|
||||||
|
json.dump(gt_entries, f, indent=2)
|
||||||
|
|
||||||
|
print(f"[prep] crops saved: {len(gt_entries)}", file=sys.stderr)
|
||||||
|
print(f"[prep] crop failures: {n_failed}", file=sys.stderr)
|
||||||
|
print(f"[prep] actors covered: {len(per_actor_count)}/{len(gallery_ids)}", file=sys.stderr)
|
||||||
|
for imdb_id, name in sorted(gallery_ids.items()):
|
||||||
|
n = per_actor_count.get(imdb_id, 0)
|
||||||
|
status = f"{n} crops" if n else "NO MATCH"
|
||||||
|
print(f" {name:30s} {status}", file=sys.stderr)
|
||||||
|
print(f"[prep] gt.json → {gt_path}", file=sys.stderr)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,144 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
movienet_score.py — compare model predictions against ground truth.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python scripts/movienet_score.py \
|
||||||
|
--gt eval/gt.json \
|
||||||
|
--predictions eval/predictions_r50.json:R50:167MB \
|
||||||
|
eval/predictions_r18.json:R18:46MB \
|
||||||
|
eval/predictions_mbf.json:MBF:13MB \
|
||||||
|
[--output eval/report.md]
|
||||||
|
|
||||||
|
Each --predictions value is <path>:<label>:<size> (size is display-only).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
from collections import defaultdict
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def load(path: str) -> list[dict]:
|
||||||
|
with open(path) as f:
|
||||||
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
def score(predictions: list[dict], gt_map: dict[str, str]) -> dict:
|
||||||
|
n_total = len(predictions)
|
||||||
|
n_det_fail = sum(1 for p in predictions if p["detection_failed"])
|
||||||
|
evaluated = [p for p in predictions if not p["detection_failed"]]
|
||||||
|
|
||||||
|
n_correct = sum(1 for p in evaluated if p["pred"] == p["gt"])
|
||||||
|
rank1 = n_correct / len(evaluated) * 100 if evaluated else 0.0
|
||||||
|
|
||||||
|
sims_correct = [p["similarity"] for p in evaluated if p["pred"] == p["gt"]]
|
||||||
|
mean_sim = sum(sims_correct) / len(sims_correct) if sims_correct else 0.0
|
||||||
|
|
||||||
|
# Per-actor recall
|
||||||
|
per_actor: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0, "name": ""})
|
||||||
|
for p in evaluated:
|
||||||
|
actor_id = p["gt"]
|
||||||
|
per_actor[actor_id]["total"] += 1
|
||||||
|
per_actor[actor_id]["name"] = gt_map.get(actor_id, actor_id)
|
||||||
|
if p["pred"] == actor_id:
|
||||||
|
per_actor[actor_id]["correct"] += 1
|
||||||
|
|
||||||
|
return {
|
||||||
|
"n_total": n_total,
|
||||||
|
"n_det_fail": n_det_fail,
|
||||||
|
"n_evaluated": len(evaluated),
|
||||||
|
"rank1": rank1,
|
||||||
|
"mean_sim_correct": mean_sim,
|
||||||
|
"per_actor": dict(per_actor),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def render_table(rows: list[dict], headers: list[str]) -> str:
|
||||||
|
col_widths = [max(len(h), max(len(str(r[h])) for r in rows)) for h in headers]
|
||||||
|
sep = "| " + " | ".join("-" * w for w in col_widths) + " |"
|
||||||
|
header = "| " + " | ".join(h.ljust(w) for h, w in zip(headers, col_widths)) + " |"
|
||||||
|
lines = [header, sep]
|
||||||
|
for r in rows:
|
||||||
|
lines.append("| " + " | ".join(str(r[h]).ljust(w) for h, w in zip(headers, col_widths)) + " |")
|
||||||
|
return "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p = argparse.ArgumentParser()
|
||||||
|
p.add_argument("--gt", required=True)
|
||||||
|
p.add_argument("--predictions", required=True, nargs="+",
|
||||||
|
metavar="PATH:LABEL:SIZE")
|
||||||
|
p.add_argument("--output", default=None)
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
gt_entries = load(args.gt)
|
||||||
|
gt_map = {e["imdb_id"]: e["actor_name"] for e in gt_entries}
|
||||||
|
|
||||||
|
models = []
|
||||||
|
for spec in args.predictions:
|
||||||
|
parts = spec.split(":")
|
||||||
|
if len(parts) != 3:
|
||||||
|
print(f"[error] expected PATH:LABEL:SIZE, got: {spec}", file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
path, label, size = parts
|
||||||
|
preds = load(path)
|
||||||
|
s = score(preds, gt_map)
|
||||||
|
models.append({"label": label, "size": size, "score": s})
|
||||||
|
|
||||||
|
# ── Summary table ────────────────────────────────────────────────────────────
|
||||||
|
summary_rows = []
|
||||||
|
for m in models:
|
||||||
|
s = m["score"]
|
||||||
|
det_fail_pct = s["n_det_fail"] / s["n_total"] * 100 if s["n_total"] else 0
|
||||||
|
summary_rows.append({
|
||||||
|
"Model": m["label"],
|
||||||
|
"Rank-1": f"{s['rank1']:.1f}%",
|
||||||
|
"Det.Fail": f"{det_fail_pct:.1f}%",
|
||||||
|
"Mean-sim": f"{s['mean_sim_correct']:.3f}",
|
||||||
|
"Probes": str(s["n_total"]),
|
||||||
|
"Size": m["size"],
|
||||||
|
})
|
||||||
|
|
||||||
|
summary_table = render_table(
|
||||||
|
summary_rows,
|
||||||
|
["Model", "Rank-1", "Det.Fail", "Mean-sim", "Probes", "Size"]
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Per-actor table (using first model's actor list as reference) ────────────
|
||||||
|
all_actor_ids = sorted({e["imdb_id"] for e in gt_entries})
|
||||||
|
actor_rows = []
|
||||||
|
for actor_id in all_actor_ids:
|
||||||
|
row = {"Actor": gt_map.get(actor_id, actor_id)}
|
||||||
|
for m in models:
|
||||||
|
pa = m["score"]["per_actor"].get(actor_id, {"correct": 0, "total": 0})
|
||||||
|
recall = pa["correct"] / pa["total"] * 100 if pa["total"] else 0.0
|
||||||
|
row[m["label"]] = f"{recall:.0f}% ({pa['correct']}/{pa['total']})"
|
||||||
|
actor_rows.append(row)
|
||||||
|
|
||||||
|
actor_headers = ["Actor"] + [m["label"] for m in models]
|
||||||
|
actor_table = render_table(actor_rows, actor_headers)
|
||||||
|
|
||||||
|
# ── Assemble report ──────────────────────────────────────────────────────────
|
||||||
|
report = f"""# MovieNet Validation Report
|
||||||
|
|
||||||
|
## Summary
|
||||||
|
|
||||||
|
{summary_table}
|
||||||
|
|
||||||
|
## Per-Actor Recall
|
||||||
|
|
||||||
|
{actor_table}
|
||||||
|
"""
|
||||||
|
|
||||||
|
print(report)
|
||||||
|
if args.output:
|
||||||
|
out = Path(args.output)
|
||||||
|
out.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
out.write_text(report)
|
||||||
|
print(f"[score] report written → {out}", file=sys.stderr)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,132 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "face_utils.hpp"
|
||||||
|
|
||||||
|
#include <onnxruntime/onnxruntime_cxx_api.h>
|
||||||
|
#include <opencv2/dnn.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// ── ArcFaceEmbedder ───────────────────────────────────────────────────────────
|
||||||
|
// ONNX Runtime-based embedder for InsightFace ArcFace (w600k_r50, mbf, r18).
|
||||||
|
// Replaces cv::dnn::Net which has no GPU path and is ~5–10× slower.
|
||||||
|
//
|
||||||
|
// Input: [N, 3, 112, 112] float32, BGR→RGB, normalised to [-1, 1]
|
||||||
|
// Output: [N, 512] float32 → L2-normalised per row
|
||||||
|
//
|
||||||
|
// ORT Run() is thread-safe; no external locking is needed.
|
||||||
|
|
||||||
|
struct ArcFaceEmbedder {
|
||||||
|
explicit ArcFaceEmbedder(const std::string& model_path,
|
||||||
|
OrtProvider provider = OrtProvider::CPU,
|
||||||
|
TrtConfig trt_cfg = {},
|
||||||
|
int max_batch = 4)
|
||||||
|
{
|
||||||
|
Ort::SessionOptions opts;
|
||||||
|
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||||
|
opts.SetIntraOpNumThreads(1);
|
||||||
|
|
||||||
|
if (provider == OrtProvider::TensorRT) {
|
||||||
|
if (trt_cfg.input_name.empty()) {
|
||||||
|
// Probe input name from a tiny CPU session so we can configure
|
||||||
|
// the dynamic-batch profile before the real session is built.
|
||||||
|
Ort::SessionOptions probe_opts;
|
||||||
|
probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
|
||||||
|
Ort::Session probe(env_, model_path.c_str(), probe_opts);
|
||||||
|
Ort::AllocatorWithDefaultOptions alloc;
|
||||||
|
trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
|
||||||
|
}
|
||||||
|
if (trt_cfg.profile_min.empty()) {
|
||||||
|
const std::string tail = "x3x112x112";
|
||||||
|
const int m = std::max(1, max_batch);
|
||||||
|
trt_cfg.profile_min = "1" + tail;
|
||||||
|
trt_cfg.profile_opt = std::to_string(m) + tail;
|
||||||
|
trt_cfg.profile_max = std::to_string(m) + tail;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
apply_ort_provider(opts, provider, "ArcFace", trt_cfg);
|
||||||
|
|
||||||
|
session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);
|
||||||
|
|
||||||
|
Ort::AllocatorWithDefaultOptions alloc;
|
||||||
|
auto in_name = session_->GetInputNameAllocated(0, alloc);
|
||||||
|
auto out_name = session_->GetOutputNameAllocated(0, alloc);
|
||||||
|
input_name_ = in_name.get();
|
||||||
|
output_name_ = out_name.get();
|
||||||
|
|
||||||
|
auto in_type = session_->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetElementType();
|
||||||
|
auto out_type = session_->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetElementType();
|
||||||
|
input_is_fp16_ = (in_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
|
||||||
|
output_is_fp16_ = (out_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
|
||||||
|
|
||||||
|
std::cerr << "[ArcFace] loaded: " << model_path << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
// Embed a batch of 112×112 BGR crops. Returns L2-normalised 512-d embeddings.
|
||||||
|
std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) const {
|
||||||
|
if (crops.empty()) return {};
|
||||||
|
const int n = static_cast<int>(crops.size());
|
||||||
|
|
||||||
|
// BGR→RGB, build NCHW float32 blob normalised to [-1, 1]
|
||||||
|
std::vector<cv::Mat> rgbs(n);
|
||||||
|
for (int i = 0; i < n; ++i)
|
||||||
|
cv::cvtColor(crops[i], rgbs[i], cv::COLOR_BGR2RGB);
|
||||||
|
|
||||||
|
cv::Mat blob = cv::dnn::blobFromImages(
|
||||||
|
rgbs, 1.0 / 128.0, {112, 112},
|
||||||
|
cv::Scalar(127.5, 127.5, 127.5),
|
||||||
|
/*swapRB=*/false, /*crop=*/false, CV_32F);
|
||||||
|
|
||||||
|
const std::array<int64_t, 4> in_shape = {n, 3, 112, 112};
|
||||||
|
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||||||
|
|
||||||
|
const char* in_name = input_name_.c_str();
|
||||||
|
const char* out_name = output_name_.c_str();
|
||||||
|
|
||||||
|
cv::Mat blob16;
|
||||||
|
if (input_is_fp16_) blob.convertTo(blob16, CV_16F);
|
||||||
|
|
||||||
|
Ort::Value in_tensor = input_is_fp16_
|
||||||
|
? Ort::Value::CreateTensor<Ort::Float16_t>(
|
||||||
|
mem, reinterpret_cast<Ort::Float16_t*>(blob16.ptr<uint16_t>()), blob16.total(),
|
||||||
|
in_shape.data(), in_shape.size())
|
||||||
|
: Ort::Value::CreateTensor<float>(
|
||||||
|
mem, blob.ptr<float>(), blob.total(),
|
||||||
|
in_shape.data(), in_shape.size());
|
||||||
|
|
||||||
|
auto outs = session_->Run(Ort::RunOptions{nullptr}, &in_name, &in_tensor, 1, &out_name, 1);
|
||||||
|
|
||||||
|
std::vector<Embedding> result(n);
|
||||||
|
if (output_is_fp16_) {
|
||||||
|
const auto* data16 = outs[0].GetTensorData<Ort::Float16_t>();
|
||||||
|
std::vector<float> buf(n * 512);
|
||||||
|
for (int j = 0; j < n * 512; ++j)
|
||||||
|
buf[j] = data16[j].ToFloat();
|
||||||
|
for (int i = 0; i < n; ++i)
|
||||||
|
result[i] = l2_normalise(buf.data() + i * 512);
|
||||||
|
} else {
|
||||||
|
const float* data = outs[0].GetTensorData<float>();
|
||||||
|
for (int i = 0; i < n; ++i)
|
||||||
|
result[i] = l2_normalise(data + i * 512);
|
||||||
|
}
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
Embedding embed_one(const cv::Mat& crop) const {
|
||||||
|
return embed({crop})[0];
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
Ort::Env env_{ORT_LOGGING_LEVEL_ERROR, "arcface"};
|
||||||
|
std::unique_ptr<Ort::Session> session_;
|
||||||
|
std::string input_name_;
|
||||||
|
std::string output_name_;
|
||||||
|
bool input_is_fp16_ = false;
|
||||||
|
bool output_is_fp16_ = false;
|
||||||
|
};
|
||||||
@@ -0,0 +1,85 @@
|
|||||||
|
// build_gallery — build an actor embedding gallery from a directory of images
|
||||||
|
//
|
||||||
|
// Gallery directory layout:
|
||||||
|
// gallery_root/
|
||||||
|
// nm0000093_Brad_Pitt/
|
||||||
|
// img1.jpg
|
||||||
|
// img2.jpg
|
||||||
|
// nm0000129_Cate_Blanchett/
|
||||||
|
// ...
|
||||||
|
//
|
||||||
|
// Usage:
|
||||||
|
// build_gallery --root <gallery_root> --output <gallery.json> [options]
|
||||||
|
//
|
||||||
|
// Options:
|
||||||
|
// --detector <path> SCRFD detector model (default: models/scrfd_500m_bnkps.onnx)
|
||||||
|
// --arcface <path> ArcFace model (default: models/arcface_w600k_r50.onnx)
|
||||||
|
// --conf <f> face detection confidence threshold (default: 0.5)
|
||||||
|
// --nms <f> NMS IoU threshold (default: 0.4)
|
||||||
|
|
||||||
|
#include "gallery/gallery_builder.hpp"
|
||||||
|
#include "gallery/gallery_store.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <cstring>
|
||||||
|
#include <iostream>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
int main(int argc, char** argv) {
|
||||||
|
std::string root_path, output_path;
|
||||||
|
std::string detector_model = kDefaultDetectorModel;
|
||||||
|
std::string arcface_model = kDefaultArcfaceModel;
|
||||||
|
float conf = 0.5f, nms_thr = 0.4f;
|
||||||
|
int max_side = 500;
|
||||||
|
|
||||||
|
for (int i = 1; i < argc; ++i) {
|
||||||
|
auto arg = [&](const char* f) { return std::strcmp(argv[i], f) == 0; };
|
||||||
|
auto next = [&]() -> std::string {
|
||||||
|
if (++i >= argc)
|
||||||
|
throw std::runtime_error(std::string("missing arg after ") + argv[i-1]);
|
||||||
|
return argv[i];
|
||||||
|
};
|
||||||
|
try {
|
||||||
|
if (arg("--root")) root_path = next();
|
||||||
|
else if (arg("--output")) output_path = next();
|
||||||
|
else if (arg("--detector")) detector_model = next();
|
||||||
|
else if (arg("--arcface")) arcface_model = next();
|
||||||
|
else if (arg("--conf")) conf = std::stof(next());
|
||||||
|
else if (arg("--nms")) nms_thr = std::stof(next());
|
||||||
|
else if (arg("--max-side")) max_side = std::stoi(next());
|
||||||
|
else { std::cerr << "[warn] unknown flag: " << argv[i] << "\n"; }
|
||||||
|
} catch (const std::exception& e) {
|
||||||
|
std::cerr << "Error: " << e.what() << "\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (root_path.empty() || output_path.empty()) {
|
||||||
|
std::cerr << "Usage: build_gallery --root <dir> --output <gallery.json> "
|
||||||
|
"[--detector <path>] [--arcface <path>] [--max-side <N>]\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
BuildConfig cfg;
|
||||||
|
cfg.gallery_root = root_path;
|
||||||
|
cfg.detector_model = detector_model;
|
||||||
|
cfg.arcface_model = arcface_model;
|
||||||
|
cfg.detector_conf = conf;
|
||||||
|
cfg.detector_nms = nms_thr;
|
||||||
|
cfg.max_side = max_side;
|
||||||
|
|
||||||
|
try {
|
||||||
|
ActorGallery gallery = build_gallery(cfg);
|
||||||
|
if (gallery.actors.empty()) {
|
||||||
|
std::cerr << "No actors built — check your gallery directory.\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
save_gallery(output_path, gallery);
|
||||||
|
std::cerr << "Gallery saved to: " << output_path << "\n";
|
||||||
|
} catch (const std::exception& e) {
|
||||||
|
std::cerr << "Fatal: " << e.what() << "\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
@@ -0,0 +1,72 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
inline const std::string kDefaultDetectorModel = std::string(SAE_MODELS_DIR) + "/scrfd_500m_bnkps.onnx";
|
||||||
|
inline const std::string kDefaultArcfaceModel = std::string(SAE_MODELS_DIR) + "/arcface_w600k_r50.onnx";
|
||||||
|
|
||||||
|
enum class Verbosity {
|
||||||
|
minimal, // actor names + merged time windows only
|
||||||
|
standard, // per-frame detail: bbox, similarity, unknowns logged
|
||||||
|
xray, // Jellyfin-Xray format: {"second": ["Actor", ...], ...}
|
||||||
|
};
|
||||||
|
// debug verbosity = compile with -DSAE_DEBUG → scene_analyze_debug binary
|
||||||
|
|
||||||
|
struct Config {
|
||||||
|
// ── Input ─────────────────────────────────────────────────────────────────
|
||||||
|
std::string movie_path;
|
||||||
|
std::string gallery_path; // gallery.json produced by build_gallery
|
||||||
|
|
||||||
|
// ── Output ───────────────────────────────────────────────────────────────
|
||||||
|
std::string output_path; // annotations.json
|
||||||
|
Verbosity verbosity{Verbosity::minimal};
|
||||||
|
|
||||||
|
// ── Sampling ─────────────────────────────────────────────────────────────
|
||||||
|
float sample_fps{1.0f}; // frames to analyse per second of movie
|
||||||
|
float max_decode_fps{0.f}; // wall-clock cap on source decode rate (0 = uncapped)
|
||||||
|
double start_sec{0.0}; // seek to this timestamp before sampling
|
||||||
|
double end_sec{-1.0}; // stop at this timestamp (-1 = end of file)
|
||||||
|
|
||||||
|
// ── Detection (SCRFD-500MF via cv::dnn::Net) ──────────────────────────────
|
||||||
|
std::string detector_model;
|
||||||
|
std::string detector_engine; // optional path to pre-built TRT engine; bypasses ORT
|
||||||
|
int max_faces{10}; // pipeline cap: keep only the N largest faces
|
||||||
|
float min_face_px{40.f}; // discard detections narrower or shorter than this
|
||||||
|
float detector_conf{0.5f};
|
||||||
|
float detector_nms{0.4f};
|
||||||
|
|
||||||
|
// ── Recognition (ArcFace ONNX) ────────────────────────────────────────────
|
||||||
|
std::string arcface_model;
|
||||||
|
std::string arcface_engine; // optional path to a pre-built TRT engine; bypasses ORT
|
||||||
|
int embed_batch_size{4}; // max faces per ORT Run() call — bounds per-call latency
|
||||||
|
float match_prior{0.5f}; // base-rate prior; 0.5 = use calibrated sigmoid directly
|
||||||
|
float prob_threshold{0.99f}; // posterior P(match | sim, prior) threshold
|
||||||
|
float match_threshold{0.45f}; // cosine distance hard ceiling fallback (no calibration)
|
||||||
|
float match_ratio{0.80f}; // ratio test fallback: accept if best/second < ratio
|
||||||
|
float match_ratio_ceil{0.65f}; // ratio test only fires below this absolute distance
|
||||||
|
|
||||||
|
// ── Cut detection ────────────────────────────────────────────────────────
|
||||||
|
float cut_threshold{0.70f}; // grayscale histogram correlation below this → hard cut
|
||||||
|
|
||||||
|
// ── Face tracking (frame-to-frame) ───────────────────────────────────────
|
||||||
|
float track_alpha{0.4f}; // cost weight: 0=embedding only, 1=spatial only
|
||||||
|
float track_min_iou{0.1f}; // IoU below which spatial link alone is rejected
|
||||||
|
float track_max_embed_dist{0.7f}; // cosine dist above which embedding link alone is rejected
|
||||||
|
int track_max_frames_missing{5}; // expire track after N consecutive missed frames
|
||||||
|
int track_min_frames{3}; // frames before track mean replaces per-frame embedding
|
||||||
|
|
||||||
|
// ── Scene tracking ────────────────────────────────────────────────────────
|
||||||
|
double extinction_sec{5.0}; // keep actor active this many seconds after last detection
|
||||||
|
double anneal_sec{2.0}; // merge actor windows separated by less than this into one epoch
|
||||||
|
|
||||||
|
// ── TensorRT ──────────────────────────────────────────────────────────────
|
||||||
|
// Only active when OrtProvider::TensorRT is detected.
|
||||||
|
// INT8 is unsafe for ArcFace without a calibration table.
|
||||||
|
TrtConfig trt{}; // fp16=true, int8=false, cache_dir="./trt_cache"
|
||||||
|
|
||||||
|
// ── Debug output (only used when SAE_DEBUG is defined) ───────────────────
|
||||||
|
#ifdef SAE_DEBUG
|
||||||
|
std::string debug_dir{"debug_frames"};
|
||||||
|
float crop_context{1.5f}; // bbox expansion factor for context crop
|
||||||
|
#endif
|
||||||
|
};
|
||||||
@@ -0,0 +1,251 @@
|
|||||||
|
// embed_faces — run SCRFD-500MF + ArcFace on a list of image files and write
|
||||||
|
// embeddings as JSON to stdout.
|
||||||
|
//
|
||||||
|
// Usage:
|
||||||
|
// embed_faces [--detector <path>] [--arcface <path>] [--conf <f>] [--nms <f>]
|
||||||
|
// image1.jpg image2.jpg ...
|
||||||
|
//
|
||||||
|
// Output (stdout): JSON array, one object per input image:
|
||||||
|
// [
|
||||||
|
// {
|
||||||
|
// "image": "actor.jpg",
|
||||||
|
// "embedding": [0.012, -0.034, ...], // 512 floats, L2-normalised
|
||||||
|
// "bbox": [x, y, w, h],
|
||||||
|
// "confidence": 0.91
|
||||||
|
// },
|
||||||
|
// {
|
||||||
|
// "image": "bad.jpg",
|
||||||
|
// "embedding": null, // no face detected / alignment failed
|
||||||
|
// "error": "no face detected"
|
||||||
|
// }
|
||||||
|
// ]
|
||||||
|
//
|
||||||
|
// Design: each image is processed independently. If multiple faces are
|
||||||
|
// detected the one with the highest confidence is used (gallery images are
|
||||||
|
// expected to contain exactly one subject). A warning is printed to stderr
|
||||||
|
// when more than one face is found.
|
||||||
|
//
|
||||||
|
// This binary is intentionally a thin wrapper around the same ONNX models
|
||||||
|
// used by scene_analyze, so embeddings are guaranteed compatible.
|
||||||
|
|
||||||
|
#include "arcface_embedder.hpp"
|
||||||
|
#include "trt_arcface_embedder.hpp"
|
||||||
|
#include "trt_scrfd_decoder.hpp"
|
||||||
|
#include "face_utils.hpp"
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "scrfd_decoder.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <opencv2/imgcodecs.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <nlohmann/json.hpp>
|
||||||
|
|
||||||
|
#include <algorithm>
|
||||||
|
#include <cstring>
|
||||||
|
#include <filesystem>
|
||||||
|
#include <functional>
|
||||||
|
#include <iostream>
|
||||||
|
#include <memory>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
namespace fs = std::filesystem;
|
||||||
|
|
||||||
|
using json = nlohmann::json;
|
||||||
|
|
||||||
|
// ── Per-image result ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
struct FaceResult {
|
||||||
|
std::string image_path;
|
||||||
|
bool ok{false};
|
||||||
|
std::string error;
|
||||||
|
Embedding embedding{};
|
||||||
|
float confidence{0.f};
|
||||||
|
float bbox[4]{}; // x, y, w, h
|
||||||
|
std::array<cv::Point2f, 5> landmarks{};
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── Debug rendering ───────────────────────────────────────────────────────────
|
||||||
|
// Writes <dir>/<stem>_annotated.jpg (input with bbox + 5 landmarks) and
|
||||||
|
// <dir>/<stem>_aligned.jpg (112×112 aligned crop). Stem is derived from the
|
||||||
|
// parent directory and filename so images from different actor folders don't
|
||||||
|
// collide when fed into a single debug dir.
|
||||||
|
|
||||||
|
static std::string debug_stem(const std::string& path) {
|
||||||
|
fs::path p(path);
|
||||||
|
std::string parent = p.parent_path().filename().string();
|
||||||
|
std::string stem = p.stem().string();
|
||||||
|
return parent.empty() ? stem : parent + "_" + stem;
|
||||||
|
}
|
||||||
|
|
||||||
|
static void save_debug(const std::string& dir,
|
||||||
|
const std::string& src_path,
|
||||||
|
const cv::Mat& img,
|
||||||
|
const DetectedFace& face,
|
||||||
|
const cv::Mat& aligned) {
|
||||||
|
fs::create_directories(dir);
|
||||||
|
cv::Mat annotated = img.clone();
|
||||||
|
cv::rectangle(annotated, face.bbox, {0, 255, 0}, 2);
|
||||||
|
static const cv::Scalar colors[5] = {
|
||||||
|
{ 0, 0, 255}, // right eye — red
|
||||||
|
{255, 0, 0}, // left eye — blue
|
||||||
|
{ 0, 255, 255}, // nose — yellow
|
||||||
|
{ 0, 255, 0}, // right mouth — green
|
||||||
|
{255, 0, 255}, // left mouth — magenta
|
||||||
|
};
|
||||||
|
for (int i = 0; i < 5; ++i)
|
||||||
|
cv::circle(annotated, face.landmarks[i], 4, colors[i], -1);
|
||||||
|
|
||||||
|
const std::string stem = debug_stem(src_path);
|
||||||
|
cv::imwrite(dir + "/" + stem + "_annotated.jpg", annotated);
|
||||||
|
cv::imwrite(dir + "/" + stem + "_aligned.jpg", aligned);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Process one image ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
static FaceResult process(const std::string& path,
|
||||||
|
const std::function<std::vector<DetectedFace>(const cv::Mat&)>& detect,
|
||||||
|
const std::function<Embedding(const cv::Mat&)>& embed_one,
|
||||||
|
int max_side,
|
||||||
|
const std::string& debug_dir = "") {
|
||||||
|
FaceResult res;
|
||||||
|
res.image_path = path;
|
||||||
|
|
||||||
|
cv::Mat img = cv::imread(path);
|
||||||
|
if (img.empty()) {
|
||||||
|
res.error = "cannot read image";
|
||||||
|
return res;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (max_side > 0) {
|
||||||
|
const int big = std::max(img.cols, img.rows);
|
||||||
|
if (big > max_side) {
|
||||||
|
const double s = static_cast<double>(max_side) / big;
|
||||||
|
cv::resize(img, img, {}, s, s, cv::INTER_AREA);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<DetectedFace> faces = detect(img);
|
||||||
|
|
||||||
|
if (faces.empty()) {
|
||||||
|
res.error = "no face detected";
|
||||||
|
return res;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (faces.size() > 1)
|
||||||
|
std::cerr << "[warn] " << path << ": " << faces.size()
|
||||||
|
<< " faces detected, using highest-confidence one\n";
|
||||||
|
|
||||||
|
const auto& best = *std::max_element(
|
||||||
|
faces.begin(), faces.end(),
|
||||||
|
[](const DetectedFace& a, const DetectedFace& b) {
|
||||||
|
return a.confidence < b.confidence;
|
||||||
|
});
|
||||||
|
|
||||||
|
cv::Mat crop = align_face(img, best.landmarks);
|
||||||
|
if (crop.empty()) {
|
||||||
|
res.error = "alignment failed";
|
||||||
|
return res;
|
||||||
|
}
|
||||||
|
|
||||||
|
res.ok = true;
|
||||||
|
res.embedding = embed_one(crop);
|
||||||
|
res.confidence = best.confidence;
|
||||||
|
res.bbox[0] = best.bbox.x;
|
||||||
|
res.bbox[1] = best.bbox.y;
|
||||||
|
res.bbox[2] = best.bbox.width;
|
||||||
|
res.bbox[3] = best.bbox.height;
|
||||||
|
res.landmarks = best.landmarks;
|
||||||
|
|
||||||
|
if (!debug_dir.empty())
|
||||||
|
save_debug(debug_dir, path, img, best, crop);
|
||||||
|
|
||||||
|
return res;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Main ──────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
int main(int argc, char** argv) {
|
||||||
|
std::string detector_model = kDefaultDetectorModel;
|
||||||
|
std::string detector_engine;
|
||||||
|
std::string arcface_model = kDefaultArcfaceModel;
|
||||||
|
std::string arcface_engine;
|
||||||
|
std::string debug_dir;
|
||||||
|
float conf = 0.5f, nms = 0.4f;
|
||||||
|
int max_side = 500;
|
||||||
|
std::vector<std::string> images;
|
||||||
|
|
||||||
|
for (int i = 1; i < argc; ++i) {
|
||||||
|
if (std::strcmp(argv[i], "--detector") == 0 && i+1 < argc) { detector_model = argv[++i]; }
|
||||||
|
else if (std::strcmp(argv[i], "--detector-engine") == 0 && i+1 < argc) { detector_engine = argv[++i]; }
|
||||||
|
else if (std::strcmp(argv[i], "--arcface") == 0 && i+1 < argc) { arcface_model = argv[++i]; }
|
||||||
|
else if (std::strcmp(argv[i], "--arcface-engine") == 0 && i+1 < argc) { arcface_engine = argv[++i]; }
|
||||||
|
else if (std::strcmp(argv[i], "--conf") == 0 && i+1 < argc) { conf = std::stof(argv[++i]); }
|
||||||
|
else if (std::strcmp(argv[i], "--nms") == 0 && i+1 < argc) { nms = std::stof(argv[++i]); }
|
||||||
|
else if (std::strcmp(argv[i], "--save-debug") == 0 && i+1 < argc) { debug_dir = argv[++i]; }
|
||||||
|
else if (std::strcmp(argv[i], "--max-side") == 0 && i+1 < argc) { max_side = std::stoi(argv[++i]); }
|
||||||
|
else if (argv[i][0] != '-') { images.push_back(argv[i]); }
|
||||||
|
else { std::cerr << "[warn] unknown flag: " << argv[i] << "\n"; }
|
||||||
|
}
|
||||||
|
|
||||||
|
if (images.empty()) {
|
||||||
|
std::cerr << "Usage: embed_faces [--detector <path>] [--arcface <path>] "
|
||||||
|
"[--save-debug <dir>] [--max-side <N>] image1.jpg ...\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
const OrtProvider provider = detect_ort_provider();
|
||||||
|
std::cerr << "[embed_faces] inference provider: " << provider_name(provider) << "\n";
|
||||||
|
|
||||||
|
std::unique_ptr<SCRFDDecoder> ort_det;
|
||||||
|
std::unique_ptr<TrtScrfdDecoder> trt_det;
|
||||||
|
std::function<std::vector<DetectedFace>(const cv::Mat&)> detect;
|
||||||
|
if (!detector_engine.empty()) {
|
||||||
|
trt_det = std::make_unique<TrtScrfdDecoder>(detector_engine, conf, nms);
|
||||||
|
detect = [&](const cv::Mat& im) { return trt_det->detect(im); };
|
||||||
|
} else {
|
||||||
|
ort_det = std::make_unique<SCRFDDecoder>(detector_model, conf, nms, provider);
|
||||||
|
detect = [&](const cv::Mat& im) { return ort_det->detect(im); };
|
||||||
|
}
|
||||||
|
|
||||||
|
std::unique_ptr<ArcFaceEmbedder> ort_emb;
|
||||||
|
std::unique_ptr<TrtArcFaceEmbedder> trt_emb;
|
||||||
|
std::function<Embedding(const cv::Mat&)> embed_one;
|
||||||
|
if (!arcface_engine.empty()) {
|
||||||
|
trt_emb = std::make_unique<TrtArcFaceEmbedder>(arcface_engine);
|
||||||
|
embed_one = [&](const cv::Mat& c) { return trt_emb->embed({c})[0]; };
|
||||||
|
} else {
|
||||||
|
ort_emb = std::make_unique<ArcFaceEmbedder>(arcface_model, provider);
|
||||||
|
embed_one = [&](const cv::Mat& c) { return ort_emb->embed_one(c); };
|
||||||
|
}
|
||||||
|
|
||||||
|
// Process images and build JSON output
|
||||||
|
json output = json::array();
|
||||||
|
|
||||||
|
for (const auto& path : images) {
|
||||||
|
std::cerr << "[embed_faces] " << path << "\n";
|
||||||
|
FaceResult res = process(path, detect, embed_one, max_side, debug_dir);
|
||||||
|
|
||||||
|
json entry;
|
||||||
|
entry["image"] = res.image_path;
|
||||||
|
if (res.ok) {
|
||||||
|
entry["embedding"] = std::vector<float>(res.embedding.begin(),
|
||||||
|
res.embedding.end());
|
||||||
|
entry["confidence"] = res.confidence;
|
||||||
|
entry["bbox"] = {res.bbox[0], res.bbox[1], res.bbox[2], res.bbox[3]};
|
||||||
|
json lms = json::array();
|
||||||
|
for (const auto& pt : res.landmarks) lms.push_back({pt.x, pt.y});
|
||||||
|
entry["landmarks"] = std::move(lms);
|
||||||
|
} else {
|
||||||
|
entry["embedding"] = nullptr;
|
||||||
|
entry["error"] = res.error;
|
||||||
|
std::cerr << " [skip] " << res.error << "\n";
|
||||||
|
}
|
||||||
|
output.push_back(std::move(entry));
|
||||||
|
}
|
||||||
|
|
||||||
|
std::cout << output.dump() << "\n";
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
|
||||||
|
#include <opencv2/calib3d.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <cmath>
|
||||||
|
|
||||||
|
// ── align_face ────────────────────────────────────────────────────────────────
|
||||||
|
// Produces a 112×112 BGR crop using the ArcFace 5-point similarity transform.
|
||||||
|
// Returns an empty Mat if the affine fit fails (degenerate detection).
|
||||||
|
inline cv::Mat align_face(const cv::Mat& img,
|
||||||
|
const std::array<cv::Point2f, 5>& landmarks) {
|
||||||
|
std::vector<cv::Point2f> src(landmarks.begin(), landmarks.end());
|
||||||
|
std::vector<cv::Point2f> dst(5);
|
||||||
|
for (int i = 0; i < 5; ++i) dst[i] = {kArcFaceRef[i][0], kArcFaceRef[i][1]};
|
||||||
|
|
||||||
|
cv::Mat M = cv::estimateAffinePartial2D(src, dst, cv::noArray(), cv::RANSAC, 3.0);
|
||||||
|
if (M.empty()) return {};
|
||||||
|
|
||||||
|
cv::Mat crop;
|
||||||
|
cv::warpAffine(img, crop, M, {112, 112},
|
||||||
|
cv::INTER_LINEAR, cv::BORDER_CONSTANT, {0, 0, 0});
|
||||||
|
return crop;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── l2_normalise ──────────────────────────────────────────────────────────────
|
||||||
|
inline Embedding l2_normalise(const float* row) {
|
||||||
|
float norm = 0.f;
|
||||||
|
for (int d = 0; d < 512; ++d) norm += row[d] * row[d];
|
||||||
|
norm = std::sqrt(norm);
|
||||||
|
if (norm < 1e-6f) norm = 1e-6f;
|
||||||
|
Embedding emb;
|
||||||
|
for (int d = 0; d < 512; ++d) emb[d] = row[d] / norm;
|
||||||
|
return emb;
|
||||||
|
}
|
||||||
@@ -0,0 +1,231 @@
|
|||||||
|
#pragma once
|
||||||
|
|
||||||
|
extern "C" {
|
||||||
|
#include <libavformat/avformat.h>
|
||||||
|
#include <libavcodec/avcodec.h>
|
||||||
|
#include <libavutil/avutil.h>
|
||||||
|
#include <libavutil/imgutils.h>
|
||||||
|
#include <libavutil/opt.h>
|
||||||
|
#include <libswscale/swscale.h>
|
||||||
|
}
|
||||||
|
|
||||||
|
#include <opencv2/core.hpp>
|
||||||
|
|
||||||
|
#include <iostream>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
// ── FFmpegDecoder ─────────────────────────────────────────────────────────────
|
||||||
|
// Seek-and-decode video reader backed by FFmpeg.
|
||||||
|
//
|
||||||
|
// Hardware decode priority: NVDEC (_cuvid variants) → CPU software.
|
||||||
|
// _cuvid decoders output NV12 to system memory directly — no explicit GPU
|
||||||
|
// frame transfer is needed. swscale converts NV12/YUV → BGR24 for the rest
|
||||||
|
// of the pipeline.
|
||||||
|
//
|
||||||
|
// Non-copyable; wrap in unique_ptr if you need to move it.
|
||||||
|
|
||||||
|
struct FFmpegDecoder {
|
||||||
|
explicit FFmpegDecoder(const std::string& path, bool use_hw = true) {
|
||||||
|
if (avformat_open_input(&fmt_ctx_, path.c_str(), nullptr, nullptr) < 0)
|
||||||
|
throw std::runtime_error("[FFmpegDecoder] cannot open: " + path);
|
||||||
|
if (avformat_find_stream_info(fmt_ctx_, nullptr) < 0)
|
||||||
|
throw std::runtime_error("[FFmpegDecoder] stream info failed");
|
||||||
|
|
||||||
|
stream_idx_ = av_find_best_stream(
|
||||||
|
fmt_ctx_, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
|
||||||
|
if (stream_idx_ < 0)
|
||||||
|
throw std::runtime_error("[FFmpegDecoder] no video stream");
|
||||||
|
|
||||||
|
AVStream* stream = fmt_ctx_->streams[stream_idx_];
|
||||||
|
AVCodecID cid = stream->codecpar->codec_id;
|
||||||
|
|
||||||
|
// Try NVDEC first; fall back to software on any failure
|
||||||
|
if (use_hw) {
|
||||||
|
if (const AVCodec* hwc = hw_codec_for(cid)) {
|
||||||
|
codec_ctx_ = avcodec_alloc_context3(hwc);
|
||||||
|
avcodec_parameters_to_context(codec_ctx_, stream->codecpar);
|
||||||
|
codec_ctx_->thread_count = 1;
|
||||||
|
|
||||||
|
AVDictionary* opts = nullptr;
|
||||||
|
av_dict_set(&opts, "gpu", "0", 0);
|
||||||
|
if (avcodec_open2(codec_ctx_, hwc, &opts) >= 0) {
|
||||||
|
hw_active_ = true;
|
||||||
|
std::cerr << "[FFmpegDecoder] " << path
|
||||||
|
<< " codec=" << hwc->name << " (NVDEC)\n";
|
||||||
|
} else {
|
||||||
|
avcodec_free_context(&codec_ctx_);
|
||||||
|
std::cerr << "[FFmpegDecoder] NVDEC init failed, falling back to CPU\n";
|
||||||
|
}
|
||||||
|
av_dict_free(&opts);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!hw_active_) {
|
||||||
|
const AVCodec* swc = avcodec_find_decoder(cid);
|
||||||
|
if (!swc) throw std::runtime_error("[FFmpegDecoder] no software decoder");
|
||||||
|
codec_ctx_ = avcodec_alloc_context3(swc);
|
||||||
|
avcodec_parameters_to_context(codec_ctx_, stream->codecpar);
|
||||||
|
codec_ctx_->thread_count = 0; // auto: use all cores
|
||||||
|
if (avcodec_open2(codec_ctx_, swc, nullptr) < 0)
|
||||||
|
throw std::runtime_error("[FFmpegDecoder] cannot open codec");
|
||||||
|
std::cerr << "[FFmpegDecoder] " << path
|
||||||
|
<< " codec=" << swc->name << " (CPU)\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
frame_ = av_frame_alloc();
|
||||||
|
tmp_frame_= av_frame_alloc();
|
||||||
|
pkt_ = av_packet_alloc();
|
||||||
|
|
||||||
|
// Decode forward (no seek) when target is within this many pts units.
|
||||||
|
// 5 seconds covers typical H.264/H.265 GOP sizes.
|
||||||
|
max_forward_pts_ = to_stream_pts(5.0);
|
||||||
|
|
||||||
|
double total = duration_sec();
|
||||||
|
double vfps = fps();
|
||||||
|
std::cerr << "[FFmpegDecoder] duration=" << total
|
||||||
|
<< "s fps=" << vfps << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
~FFmpegDecoder() {
|
||||||
|
if (sws_ctx_) sws_freeContext(sws_ctx_);
|
||||||
|
av_frame_free(&frame_);
|
||||||
|
av_frame_free(&tmp_frame_);
|
||||||
|
av_packet_free(&pkt_);
|
||||||
|
avcodec_free_context(&codec_ctx_);
|
||||||
|
avformat_close_input(&fmt_ctx_);
|
||||||
|
}
|
||||||
|
|
||||||
|
FFmpegDecoder(const FFmpegDecoder&) = delete;
|
||||||
|
FFmpegDecoder& operator=(const FFmpegDecoder&) = delete;
|
||||||
|
|
||||||
|
double duration_sec() const {
|
||||||
|
if (!fmt_ctx_ || fmt_ctx_->duration == AV_NOPTS_VALUE) return 0.0;
|
||||||
|
return static_cast<double>(fmt_ctx_->duration) / AV_TIME_BASE;
|
||||||
|
}
|
||||||
|
|
||||||
|
double fps() const {
|
||||||
|
AVStream* s = fmt_ctx_->streams[stream_idx_];
|
||||||
|
if (s->avg_frame_rate.den == 0) return 25.0;
|
||||||
|
return av_q2d(s->avg_frame_rate);
|
||||||
|
}
|
||||||
|
|
||||||
|
bool hw_active() const { return hw_active_; }
|
||||||
|
const char* codec_name() const { return codec_ctx_ ? codec_ctx_->codec->name : "unknown"; }
|
||||||
|
|
||||||
|
// Decode the frame at target_sec and return it as BGR cv::Mat.
|
||||||
|
// Returns an empty Mat at EOF.
|
||||||
|
//
|
||||||
|
// Smart seek: if the target is within max_forward_sec_ ahead of the last
|
||||||
|
// decoded position, decode forward (no seek, no flush). This is dramatically
|
||||||
|
// faster for sequential sampling because avcodec_flush_buffers + re-init on
|
||||||
|
// every call is the main bottleneck — especially with NVDEC.
|
||||||
|
cv::Mat read_at(double target_sec) {
|
||||||
|
AVStream* stream = fmt_ctx_->streams[stream_idx_];
|
||||||
|
int64_t tgt_pts = to_stream_pts(target_sec);
|
||||||
|
|
||||||
|
// Decide: seek or decode forward?
|
||||||
|
bool need_seek = (last_pts_ == AV_NOPTS_VALUE) ||
|
||||||
|
(tgt_pts < last_pts_) ||
|
||||||
|
(tgt_pts - last_pts_ > max_forward_pts_);
|
||||||
|
|
||||||
|
if (need_seek) {
|
||||||
|
if (av_seek_frame(fmt_ctx_, stream_idx_, tgt_pts, AVSEEK_FLAG_BACKWARD) < 0)
|
||||||
|
av_seek_frame(fmt_ctx_, -1,
|
||||||
|
static_cast<int64_t>(target_sec * AV_TIME_BASE),
|
||||||
|
AVSEEK_FLAG_BACKWARD);
|
||||||
|
avcodec_flush_buffers(codec_ctx_);
|
||||||
|
last_pts_ = AV_NOPTS_VALUE;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Decode forward until we reach or pass target_pts.
|
||||||
|
// Convert to BGR and unref the AVFrame immediately so NVDEC surfaces
|
||||||
|
// are returned to the pool — holding them causes surface exhaustion
|
||||||
|
// at higher sample rates.
|
||||||
|
cv::Mat out;
|
||||||
|
while (out.empty()) {
|
||||||
|
int ret = av_read_frame(fmt_ctx_, pkt_);
|
||||||
|
if (ret == AVERROR_EOF || ret < 0) break;
|
||||||
|
|
||||||
|
if (pkt_->stream_index != stream_idx_) {
|
||||||
|
av_packet_unref(pkt_);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
avcodec_send_packet(codec_ctx_, pkt_);
|
||||||
|
av_packet_unref(pkt_);
|
||||||
|
|
||||||
|
while (avcodec_receive_frame(codec_ctx_, frame_) == 0) {
|
||||||
|
int64_t pts = frame_->best_effort_timestamp;
|
||||||
|
if (pts == AV_NOPTS_VALUE) pts = frame_->pts;
|
||||||
|
last_pts_ = pts;
|
||||||
|
if (pts >= tgt_pts)
|
||||||
|
out = to_bgr(frame_);
|
||||||
|
av_frame_unref(frame_); // release NVDEC surface immediately
|
||||||
|
if (!out.empty()) break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
AVFormatContext* fmt_ctx_ = nullptr;
|
||||||
|
AVCodecContext* codec_ctx_ = nullptr;
|
||||||
|
AVFrame* frame_ = nullptr;
|
||||||
|
AVFrame* tmp_frame_ = nullptr;
|
||||||
|
AVPacket* pkt_ = nullptr;
|
||||||
|
SwsContext* sws_ctx_ = nullptr;
|
||||||
|
int stream_idx_ = -1;
|
||||||
|
bool hw_active_ = false;
|
||||||
|
int64_t last_pts_ = AV_NOPTS_VALUE;
|
||||||
|
int64_t max_forward_pts_ = AV_NOPTS_VALUE; // set after codec opens
|
||||||
|
|
||||||
|
int64_t to_stream_pts(double sec) const {
|
||||||
|
AVStream* s = fmt_ctx_->streams[stream_idx_];
|
||||||
|
return av_rescale_q(static_cast<int64_t>(sec * AV_TIME_BASE),
|
||||||
|
AV_TIME_BASE_Q, s->time_base);
|
||||||
|
}
|
||||||
|
|
||||||
|
static const AVCodec* hw_codec_for(AVCodecID id) {
|
||||||
|
const char* name = nullptr;
|
||||||
|
switch (id) {
|
||||||
|
case AV_CODEC_ID_H264: name = "h264_cuvid"; break;
|
||||||
|
case AV_CODEC_ID_HEVC: name = "hevc_cuvid"; break;
|
||||||
|
case AV_CODEC_ID_AV1: name = "av1_cuvid"; break;
|
||||||
|
case AV_CODEC_ID_MPEG2VIDEO: name = "mpeg2_cuvid"; break;
|
||||||
|
case AV_CODEC_ID_MPEG4: name = "mpeg4_cuvid"; break;
|
||||||
|
case AV_CODEC_ID_VC1: name = "vc1_cuvid"; break;
|
||||||
|
default: return nullptr;
|
||||||
|
}
|
||||||
|
return avcodec_find_decoder_by_name(name);
|
||||||
|
}
|
||||||
|
|
||||||
|
cv::Mat to_bgr(AVFrame* src) {
|
||||||
|
// _cuvid decoders output NV12 to system memory.
|
||||||
|
// Generic hwaccel would output AV_PIX_FMT_CUDA and need a transfer.
|
||||||
|
AVFrame* sw = src;
|
||||||
|
if (src->format == AV_PIX_FMT_CUDA) {
|
||||||
|
tmp_frame_->format = AV_PIX_FMT_NV12;
|
||||||
|
if (av_hwframe_transfer_data(tmp_frame_, src, 0) < 0) return {};
|
||||||
|
av_frame_copy_props(tmp_frame_, src);
|
||||||
|
sw = tmp_frame_;
|
||||||
|
}
|
||||||
|
|
||||||
|
const int w = sw->width;
|
||||||
|
const int h = sw->height;
|
||||||
|
|
||||||
|
sws_ctx_ = sws_getCachedContext(sws_ctx_,
|
||||||
|
w, h, static_cast<AVPixelFormat>(sw->format),
|
||||||
|
w, h, AV_PIX_FMT_BGR24,
|
||||||
|
SWS_BILINEAR, nullptr, nullptr, nullptr);
|
||||||
|
if (!sws_ctx_) return {};
|
||||||
|
|
||||||
|
cv::Mat out(h, w, CV_8UC3);
|
||||||
|
uint8_t* dst_data[1] = { out.data };
|
||||||
|
int dst_linesize[1] = { static_cast<int>(out.step) };
|
||||||
|
sws_scale(sws_ctx_,
|
||||||
|
sw->data, sw->linesize, 0, h,
|
||||||
|
dst_data, dst_linesize);
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
};
|
||||||
@@ -0,0 +1,115 @@
|
|||||||
|
#include "gallery_builder.hpp"
|
||||||
|
#include "arcface_embedder.hpp"
|
||||||
|
#include "face_utils.hpp"
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "scrfd_decoder.hpp"
|
||||||
|
|
||||||
|
#include <opencv2/imgcodecs.hpp>
|
||||||
|
|
||||||
|
#include <algorithm>
|
||||||
|
#include <filesystem>
|
||||||
|
#include <iostream>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
namespace fs = std::filesystem;
|
||||||
|
|
||||||
|
// ── Parse "nm0000093_Brad_Pitt" → ("nm0000093", "Brad Pitt") ─────────────────
|
||||||
|
static std::pair<std::string, std::string> parse_dir_name(const std::string& dirname) {
|
||||||
|
auto pos = dirname.find('_');
|
||||||
|
if (pos == std::string::npos) return {dirname, dirname};
|
||||||
|
|
||||||
|
std::string imdb_id = dirname.substr(0, pos);
|
||||||
|
std::string raw = dirname.substr(pos + 1);
|
||||||
|
std::string name;
|
||||||
|
name.reserve(raw.size());
|
||||||
|
for (char c : raw)
|
||||||
|
name += (c == '_' ? ' ' : c);
|
||||||
|
return {imdb_id, name};
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Public API ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
ActorGallery build_gallery(const BuildConfig& cfg) {
|
||||||
|
const OrtProvider provider = detect_ort_provider();
|
||||||
|
std::cerr << "[build_gallery] inference provider: " << provider_name(provider) << "\n";
|
||||||
|
|
||||||
|
SCRFDDecoder decoder(cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider);
|
||||||
|
ArcFaceEmbedder arcface(cfg.arcface_model, provider);
|
||||||
|
|
||||||
|
ActorGallery gallery;
|
||||||
|
|
||||||
|
for (const auto& actor_dir : fs::directory_iterator(cfg.gallery_root)) {
|
||||||
|
if (!actor_dir.is_directory()) continue;
|
||||||
|
|
||||||
|
auto [imdb_id, name] = parse_dir_name(actor_dir.path().filename().string());
|
||||||
|
std::cerr << "[build_gallery] " << name << " (" << imdb_id << ")\n";
|
||||||
|
|
||||||
|
ActorGallery::Actor actor;
|
||||||
|
actor.imdb_id = imdb_id;
|
||||||
|
actor.name = name;
|
||||||
|
|
||||||
|
static const std::vector<std::string> kExts{".jpg", ".jpeg", ".png", ".webp"};
|
||||||
|
for (const auto& img_file : fs::directory_iterator(actor_dir.path())) {
|
||||||
|
if (!img_file.is_regular_file()) continue;
|
||||||
|
std::string ext = img_file.path().extension().string();
|
||||||
|
std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
|
||||||
|
if (std::find(kExts.begin(), kExts.end(), ext) == kExts.end()) continue;
|
||||||
|
|
||||||
|
cv::Mat img = cv::imread(img_file.path().string());
|
||||||
|
if (img.empty()) {
|
||||||
|
std::cerr << " [skip] cannot read " << img_file.path().filename() << "\n";
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cfg.max_side > 0) {
|
||||||
|
const int big = std::max(img.cols, img.rows);
|
||||||
|
if (big > cfg.max_side) {
|
||||||
|
const double s = static_cast<double>(cfg.max_side) / big;
|
||||||
|
cv::resize(img, img, {}, s, s, cv::INTER_AREA);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
auto faces = decoder.detect(img);
|
||||||
|
|
||||||
|
if (faces.empty()) {
|
||||||
|
std::cerr << " [skip] no face: " << img_file.path().filename() << "\n";
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if (faces.size() > 1) {
|
||||||
|
std::cerr << " [warn] " << faces.size() << " faces, using highest confidence: "
|
||||||
|
<< img_file.path().filename() << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
const auto& best = *std::max_element(
|
||||||
|
faces.begin(), faces.end(),
|
||||||
|
[](const DetectedFace& a, const DetectedFace& b) {
|
||||||
|
return a.confidence < b.confidence;
|
||||||
|
});
|
||||||
|
|
||||||
|
cv::Mat crop = align_face(img, best.landmarks);
|
||||||
|
if (crop.empty()) {
|
||||||
|
std::cerr << " [skip] alignment failed: " << img_file.path().filename() << "\n";
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
Embedding emb = arcface.embed_one(crop);
|
||||||
|
actor.embeddings.push_back(emb);
|
||||||
|
actor.source_images.push_back(img_file.path().filename().string());
|
||||||
|
|
||||||
|
std::cerr << " [ok] " << img_file.path().filename()
|
||||||
|
<< " conf=" << best.confidence << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
if (actor.embeddings.empty()) {
|
||||||
|
std::cerr << " [warn] no valid embeddings for " << name << " — skipped\n";
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
std::cerr << " → " << actor.embeddings.size() << " embeddings\n";
|
||||||
|
gallery.actors.push_back(std::move(actor));
|
||||||
|
}
|
||||||
|
|
||||||
|
std::cerr << "[build_gallery] total: " << gallery.actors.size() << " actors\n";
|
||||||
|
return gallery;
|
||||||
|
}
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
// Build an ActorGallery from a directory tree:
|
||||||
|
//
|
||||||
|
// gallery_root/
|
||||||
|
// nm0000093_Brad_Pitt/
|
||||||
|
// img1.jpg
|
||||||
|
// img2.jpg
|
||||||
|
// ...
|
||||||
|
// nm0000129_Cate_Blanchett/
|
||||||
|
// ...
|
||||||
|
//
|
||||||
|
// Each subdirectory name is parsed as "<imdb_id>_<Name_With_Underscores>".
|
||||||
|
// For every image:
|
||||||
|
// 1. Detect face with SCRFD-500MF (expect exactly one; warn and skip if 0 or >1).
|
||||||
|
// 2. Align with ArcFace 5-point transform → 112×112 crop.
|
||||||
|
// 3. Embed with ArcFace ONNX → 512-dim L2-normalised embedding.
|
||||||
|
// All embeddings are stored (best-of-N match at query time).
|
||||||
|
//
|
||||||
|
// Returns a gallery ready to pass to save_gallery() / IdentityMatcherFunc.
|
||||||
|
|
||||||
|
struct BuildConfig {
|
||||||
|
std::string gallery_root; // directory tree described above
|
||||||
|
std::string detector_model;
|
||||||
|
std::string arcface_model;
|
||||||
|
float detector_conf{0.5f};
|
||||||
|
float detector_nms{0.4f};
|
||||||
|
int max_side{500}; // downscale source images to this max dimension
|
||||||
|
// before detection — TMDB portraits are ~2k px,
|
||||||
|
// SCRFD trains on smaller faces and detection
|
||||||
|
// confidence drops on huge inputs. 0 = disabled.
|
||||||
|
};
|
||||||
|
|
||||||
|
ActorGallery build_gallery(const BuildConfig& cfg);
|
||||||
@@ -0,0 +1,118 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
|
||||||
|
#include <cmath>
|
||||||
|
#include <iostream>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// ── GalleryCalibration ────────────────────────────────────────────────────────
|
||||||
|
// Platt-style sigmoid calibration: P(match) = σ(a · similarity + b)
|
||||||
|
// where similarity = cosine similarity ∈ [-1, 1] (dot product of L2-normalised
|
||||||
|
// ArcFace embeddings).
|
||||||
|
//
|
||||||
|
// Fitted from intra-class (positive) and inter-class (negative) pairs built
|
||||||
|
// from the gallery reference embeddings. When calibration is invalid (too few
|
||||||
|
// positive pairs), the caller should fall back to the raw threshold.
|
||||||
|
struct GalleryCalibration {
|
||||||
|
float a{10.f}; // scale (positive → higher similarity → higher probability)
|
||||||
|
float b{-5.f}; // bias (decision boundary at similarity = -b/a)
|
||||||
|
bool valid{false};
|
||||||
|
|
||||||
|
// P(match | sim) using the balanced-prior calibration.
|
||||||
|
// Pass log_prior_odds = log(p0/(1-p0)) to adjust for a real base-rate prior p0:
|
||||||
|
// P(match | sim, p0) = σ(a·sim + b + log(p0/(1-p0)))
|
||||||
|
float probability(float similarity, float log_prior_odds = 0.f) const {
|
||||||
|
float z = a * similarity + b + log_prior_odds;
|
||||||
|
if (z >= 0.f) return 1.f / (1.f + std::exp(-z));
|
||||||
|
float e = std::exp(z);
|
||||||
|
return e / (1.f + e);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Similarity at which P(match, prior) == p
|
||||||
|
float boundary_at(float p = 0.5f, float log_prior_odds = 0.f) const {
|
||||||
|
return (std::log(p / (1.f - p)) - b - log_prior_odds) / a;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Fit a logistic sigmoid to gallery pair similarities.
|
||||||
|
// Positive pairs: same actor, different reference images.
|
||||||
|
// Negative pairs: different actors (all cross-actor embedding pairs).
|
||||||
|
// Class weights balance the (typically skewed) pos/neg ratio.
|
||||||
|
// Requires ≥2 positive pairs and ≥1 negative pair.
|
||||||
|
inline GalleryCalibration calibrate_gallery(
|
||||||
|
const std::vector<Embedding>& flat_emb,
|
||||||
|
const std::vector<int>& flat_actor)
|
||||||
|
{
|
||||||
|
const int n = static_cast<int>(flat_emb.size());
|
||||||
|
|
||||||
|
std::vector<float> X; // cosine similarities
|
||||||
|
std::vector<float> Y; // labels: 1 = same actor, 0 = different
|
||||||
|
|
||||||
|
for (int i = 0; i < n; ++i) {
|
||||||
|
for (int j = i + 1; j < n; ++j) {
|
||||||
|
float dot = 0.f;
|
||||||
|
for (int k = 0; k < 512; ++k)
|
||||||
|
dot += flat_emb[i][k] * flat_emb[j][k];
|
||||||
|
X.push_back(dot);
|
||||||
|
Y.push_back(flat_actor[i] == flat_actor[j] ? 1.f : 0.f);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
int n_pos = 0;
|
||||||
|
for (float y : Y) if (y > 0.5f) ++n_pos;
|
||||||
|
int n_neg = static_cast<int>(Y.size()) - n_pos;
|
||||||
|
|
||||||
|
if (n_pos < 2 || n_neg < 1) {
|
||||||
|
std::cerr << "[calibration] insufficient pairs (+" << n_pos
|
||||||
|
<< "/-" << n_neg << ") — calibration skipped\n";
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Class weights to handle pos/neg imbalance
|
||||||
|
float total = static_cast<float>(Y.size());
|
||||||
|
float w_pos = total / (2.f * n_pos);
|
||||||
|
float w_neg = total / (2.f * n_neg);
|
||||||
|
|
||||||
|
// Gradient descent logistic regression (2 parameters: a, b)
|
||||||
|
float a = 10.f, b = -5.f;
|
||||||
|
constexpr float lr = 0.05f;
|
||||||
|
constexpr int max_iter = 20000;
|
||||||
|
constexpr float tol = 1e-7f;
|
||||||
|
|
||||||
|
for (int iter = 0; iter < max_iter; ++iter) {
|
||||||
|
float da = 0.f, db = 0.f;
|
||||||
|
for (int i = 0; i < static_cast<int>(X.size()); ++i) {
|
||||||
|
float z = a * X[i] + b;
|
||||||
|
float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z))
|
||||||
|
: std::exp(z) / (1.f + std::exp(z));
|
||||||
|
float err = sig - Y[i];
|
||||||
|
float w = (Y[i] > 0.5f) ? w_pos : w_neg;
|
||||||
|
da += w * err * X[i];
|
||||||
|
db += w * err;
|
||||||
|
}
|
||||||
|
da /= total;
|
||||||
|
db /= total;
|
||||||
|
a -= lr * da;
|
||||||
|
b -= lr * db;
|
||||||
|
if (da * da + db * db < tol * tol) break;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Training accuracy at P=0.5 decision boundary
|
||||||
|
int correct = 0;
|
||||||
|
for (int i = 0; i < static_cast<int>(X.size()); ++i) {
|
||||||
|
float z = a * X[i] + b;
|
||||||
|
float sig = (z >= 0.f) ? 1.f / (1.f + std::exp(-z))
|
||||||
|
: std::exp(z) / (1.f + std::exp(z));
|
||||||
|
if ((sig > 0.5f) == (Y[i] > 0.5f)) ++correct;
|
||||||
|
}
|
||||||
|
float acc = 100.f * correct / static_cast<float>(Y.size());
|
||||||
|
|
||||||
|
GalleryCalibration cal{a, b, true};
|
||||||
|
std::cerr << "[calibration] sigmoid fitted:"
|
||||||
|
<< " a=" << a << " b=" << b
|
||||||
|
<< " boundary(P=0.5)=sim" << cal.boundary_at(0.5f)
|
||||||
|
<< " pairs=" << Y.size()
|
||||||
|
<< " (+" << n_pos << "/-" << n_neg << ")"
|
||||||
|
<< " train_acc=" << acc << "%\n";
|
||||||
|
return cal;
|
||||||
|
}
|
||||||
@@ -0,0 +1,58 @@
|
|||||||
|
#include "gallery_store.hpp"
|
||||||
|
|
||||||
|
#include <nlohmann/json.hpp>
|
||||||
|
#include <fstream>
|
||||||
|
#include <stdexcept>
|
||||||
|
|
||||||
|
using json = nlohmann::json;
|
||||||
|
|
||||||
|
ActorGallery load_gallery(const std::string& path) {
|
||||||
|
std::ifstream f(path);
|
||||||
|
if (!f.is_open())
|
||||||
|
throw std::runtime_error("load_gallery: cannot open " + path);
|
||||||
|
|
||||||
|
json j;
|
||||||
|
f >> j;
|
||||||
|
|
||||||
|
ActorGallery gallery;
|
||||||
|
for (const auto& ja : j.at("actors")) {
|
||||||
|
ActorGallery::Actor actor;
|
||||||
|
actor.imdb_id = ja.at("imdb_id").get<std::string>();
|
||||||
|
actor.name = ja.at("name").get<std::string>();
|
||||||
|
|
||||||
|
if (ja.contains("source_images"))
|
||||||
|
actor.source_images = ja.at("source_images").get<std::vector<std::string>>();
|
||||||
|
|
||||||
|
for (const auto& je : ja.at("embeddings")) {
|
||||||
|
Embedding emb = je.get<Embedding>();
|
||||||
|
actor.embeddings.push_back(emb);
|
||||||
|
}
|
||||||
|
|
||||||
|
gallery.actors.push_back(std::move(actor));
|
||||||
|
}
|
||||||
|
return gallery;
|
||||||
|
}
|
||||||
|
|
||||||
|
void save_gallery(const std::string& path, const ActorGallery& gallery) {
|
||||||
|
json j;
|
||||||
|
j["actors"] = json::array();
|
||||||
|
|
||||||
|
for (const auto& actor : gallery.actors) {
|
||||||
|
json ja;
|
||||||
|
ja["imdb_id"] = actor.imdb_id;
|
||||||
|
ja["name"] = actor.name;
|
||||||
|
ja["source_images"] = actor.source_images;
|
||||||
|
|
||||||
|
ja["embeddings"] = json::array();
|
||||||
|
for (const auto& emb : actor.embeddings) {
|
||||||
|
ja["embeddings"].push_back(
|
||||||
|
std::vector<float>(emb.begin(), emb.end()));
|
||||||
|
}
|
||||||
|
j["actors"].push_back(std::move(ja));
|
||||||
|
}
|
||||||
|
|
||||||
|
std::ofstream f(path);
|
||||||
|
if (!f.is_open())
|
||||||
|
throw std::runtime_error("save_gallery: cannot write " + path);
|
||||||
|
f << j.dump(2) << "\n";
|
||||||
|
}
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
// Load/save the actor gallery from/to a JSON file.
|
||||||
|
//
|
||||||
|
// JSON format:
|
||||||
|
// {
|
||||||
|
// "actors": [
|
||||||
|
// {
|
||||||
|
// "imdb_id": "nm0000093",
|
||||||
|
// "name": "Brad Pitt",
|
||||||
|
// "source_images": ["img1.jpg", "img2.jpg"],
|
||||||
|
// "embeddings": [[0.012, -0.034, ...], ...] // one 512-float array per image
|
||||||
|
// }
|
||||||
|
// ]
|
||||||
|
// }
|
||||||
|
|
||||||
|
ActorGallery load_gallery(const std::string& path);
|
||||||
|
void save_gallery(const std::string& path, const ActorGallery& gallery);
|
||||||
+213
@@ -0,0 +1,213 @@
|
|||||||
|
// scene_analyze — identify actors in a movie using a KPN pipeline
|
||||||
|
//
|
||||||
|
// KPN topology (release build):
|
||||||
|
//
|
||||||
|
// [frame_source] ──Frame──► [face_detector] ──SceneFrame──► [face_aligner]
|
||||||
|
// ──AlignedSceneFrame──► [embedder] ──EmbeddedSceneFrame──►
|
||||||
|
// [identity_matcher] ──MatchedSceneFrame──► [scene_tracker]
|
||||||
|
// ──SceneAnnotation──► [result_sink]
|
||||||
|
//
|
||||||
|
// Debug build (SAE_DEBUG=1):
|
||||||
|
// [identity_matcher] output fans out to both [scene_tracker] AND [debug_renderer].
|
||||||
|
// FanoutNode<MatchedSceneFrame, 2> is auto-inserted by make_network().
|
||||||
|
//
|
||||||
|
// Usage:
|
||||||
|
// scene_analyze --movie <path> --gallery <gallery.json> [options]
|
||||||
|
//
|
||||||
|
// Options:
|
||||||
|
// --output <path> output JSON (default: annotations.json)
|
||||||
|
// --fps <N> sample rate in frames/sec (default: 1.0)
|
||||||
|
// --verbosity <0|1|2> 0=minimal, 1=standard, 2=jellyfin-xray (default: 0)
|
||||||
|
// --match-threshold <f> cosine dist threshold (default: 0.45)
|
||||||
|
// --extinction <f> actor extinction window in seconds (default: 5.0)
|
||||||
|
// --detector <path> override SCRFD detector model path
|
||||||
|
// --arcface <path> override ArcFace model path
|
||||||
|
// --max-faces <N> max faces kept per frame (default: 10)
|
||||||
|
// (SAE_DEBUG only)
|
||||||
|
// --debug-dir <path> debug frames output dir (default: debug_frames)
|
||||||
|
// --crop-context <f> bbox expansion factor for context crops (default: 1.5)
|
||||||
|
|
||||||
|
#include "config.hpp"
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "gallery/gallery_store.hpp"
|
||||||
|
#include "nodes/frame_source_node.hpp"
|
||||||
|
#include "nodes/face_detector_node.hpp"
|
||||||
|
#include "nodes/face_aligner_node.hpp"
|
||||||
|
#include "nodes/embedder_node.hpp"
|
||||||
|
#include "nodes/face_tracker_node.hpp"
|
||||||
|
#include "nodes/identity_matcher_node.hpp"
|
||||||
|
#include "nodes/scene_tracker_node.hpp"
|
||||||
|
#include "nodes/result_sink_node.hpp"
|
||||||
|
#ifdef SAE_DEBUG
|
||||||
|
#include "nodes/debug_renderer_node.hpp"
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#include <kpn/kpn.hpp>
|
||||||
|
|
||||||
|
#include <atomic>
|
||||||
|
#include <chrono>
|
||||||
|
#include <cstring>
|
||||||
|
#include <iostream>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <thread>
|
||||||
|
|
||||||
|
// ── CLI parsing ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
static Config parse_args(int argc, char** argv) {
|
||||||
|
Config cfg;
|
||||||
|
cfg.detector_model = kDefaultDetectorModel;
|
||||||
|
cfg.arcface_model = kDefaultArcfaceModel;
|
||||||
|
cfg.output_path = "annotations.json";
|
||||||
|
|
||||||
|
for (int i = 1; i < argc; ++i) {
|
||||||
|
auto arg = [&](const char* flag) { return std::strcmp(argv[i], flag) == 0; };
|
||||||
|
auto next = [&]() -> std::string {
|
||||||
|
if (++i >= argc) throw std::runtime_error(std::string("missing arg after ") + argv[i-1]);
|
||||||
|
return argv[i];
|
||||||
|
};
|
||||||
|
|
||||||
|
if (arg("--movie")) cfg.movie_path = next();
|
||||||
|
else if (arg("--gallery")) cfg.gallery_path = next();
|
||||||
|
else if (arg("--output")) cfg.output_path = next();
|
||||||
|
else if (arg("--fps")) cfg.sample_fps = std::stof(next());
|
||||||
|
else if (arg("--max-decode-fps")) cfg.max_decode_fps = std::stof(next());
|
||||||
|
else if (arg("--start")) cfg.start_sec = std::stod(next());
|
||||||
|
else if (arg("--end")) cfg.end_sec = std::stod(next());
|
||||||
|
else if (arg("--cut-threshold")) cfg.cut_threshold = std::stof(next());
|
||||||
|
else if (arg("--verbosity")) { int v = std::stoi(next()); cfg.verbosity = v == 2 ? Verbosity::xray : v == 1 ? Verbosity::standard : Verbosity::minimal; }
|
||||||
|
else if (arg("--prior")) cfg.match_prior = std::stof(next());
|
||||||
|
else if (arg("--prob-threshold")) cfg.prob_threshold = std::stof(next());
|
||||||
|
else if (arg("--match-threshold")) cfg.match_threshold = std::stof(next());
|
||||||
|
else if (arg("--extinction")) cfg.extinction_sec = std::stod(next());
|
||||||
|
else if (arg("--detector")) cfg.detector_model = next();
|
||||||
|
else if (arg("--detector-engine")) cfg.detector_engine = next();
|
||||||
|
else if (arg("--arcface")) cfg.arcface_model = next();
|
||||||
|
else if (arg("--arcface-engine")) cfg.arcface_engine = next();
|
||||||
|
else if (arg("--conf")) cfg.detector_conf = std::stof(next());
|
||||||
|
else if (arg("--max-faces")) cfg.max_faces = std::stoi(next());
|
||||||
|
else if (arg("--min-face-px")) cfg.min_face_px = std::stof(next());
|
||||||
|
else if (arg("--ratio")) cfg.match_ratio = std::stof(next());
|
||||||
|
else if (arg("--ratio-ceil")) cfg.match_ratio_ceil = std::stof(next());
|
||||||
|
else if (arg("--track-alpha")) cfg.track_alpha = std::stof(next());
|
||||||
|
else if (arg("--track-min-iou")) cfg.track_min_iou = std::stof(next());
|
||||||
|
else if (arg("--track-max-embed")) cfg.track_max_embed_dist = std::stof(next());
|
||||||
|
else if (arg("--track-max-missing")) cfg.track_max_frames_missing = std::stoi(next());
|
||||||
|
else if (arg("--track-min-frames")) cfg.track_min_frames = std::stoi(next());
|
||||||
|
else if (arg("--anneal")) cfg.anneal_sec = std::stod(next());
|
||||||
|
else if (arg("--trt-cache")) cfg.trt.cache_dir = next();
|
||||||
|
else if (arg("--trt-fp16")) cfg.trt.fp16 = true;
|
||||||
|
else if (arg("--no-trt-fp16")) cfg.trt.fp16 = false;
|
||||||
|
else if (arg("--trt-int8")) cfg.trt.int8 = true;
|
||||||
|
else if (arg("--embed-batch")) cfg.embed_batch_size = std::stoi(next());
|
||||||
|
#ifdef SAE_DEBUG
|
||||||
|
else if (arg("--debug-dir")) cfg.debug_dir = next();
|
||||||
|
else if (arg("--crop-context")) cfg.crop_context = std::stof(next());
|
||||||
|
#endif
|
||||||
|
else {
|
||||||
|
std::cerr << "[warn] unknown flag: " << argv[i] << "\n";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cfg.movie_path.empty()) throw std::runtime_error("--movie is required");
|
||||||
|
if (cfg.gallery_path.empty()) throw std::runtime_error("--gallery is required");
|
||||||
|
return cfg;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Main ──────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
int main(int argc, char** argv) {
|
||||||
|
Config cfg;
|
||||||
|
try {
|
||||||
|
cfg = parse_args(argc, argv);
|
||||||
|
} catch (const std::exception& e) {
|
||||||
|
std::cerr << "Usage error: " << e.what() << "\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Load actor gallery
|
||||||
|
ActorGallery gallery;
|
||||||
|
try {
|
||||||
|
gallery = load_gallery(cfg.gallery_path);
|
||||||
|
} catch (const std::exception& e) {
|
||||||
|
std::cerr << "Gallery error: " << e.what() << "\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
std::cerr << "[main] gallery loaded: " << gallery.actors.size() << " actors\n";
|
||||||
|
|
||||||
|
// ── Construct node functors ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
std::atomic<bool> done{false};
|
||||||
|
|
||||||
|
const OrtProvider provider = detect_ort_provider();
|
||||||
|
std::cerr << "[main] inference provider: " << provider_name(provider) << "\n";
|
||||||
|
|
||||||
|
FrameSourceFunc source_fn {cfg};
|
||||||
|
FaceDetectorFunc detector_fn{cfg, provider};
|
||||||
|
FaceAlignerFunc aligner_fn;
|
||||||
|
EmbedderFunc embedder_fn{cfg, provider};
|
||||||
|
FaceTrackerFunc ftracker_fn{cfg};
|
||||||
|
IdentityMatcherFunc matcher_fn {gallery, cfg};
|
||||||
|
SceneTrackerFunc tracker_fn {cfg};
|
||||||
|
ResultSinkFunc sink_fn {cfg, done};
|
||||||
|
#ifdef SAE_DEBUG
|
||||||
|
DebugRendererFunc debug_fn {cfg};
|
||||||
|
#endif
|
||||||
|
|
||||||
|
// ── Wrap in KPN ObjectNodes ───────────────────────────────────────────────
|
||||||
|
|
||||||
|
// Queue sizes tuned to the pipeline's speed profile:
|
||||||
|
// embedder (16ms) is the slowest GPU node — buffer before it must be largest
|
||||||
|
// to prevent face_aligner pool overflows and frame drops.
|
||||||
|
kpn::ObjectNode<FrameSourceFunc, kpn::in<>, kpn::out<"frame">, "frame_source", 0> source (source_fn, 32);
|
||||||
|
kpn::ObjectNode<FaceDetectorFunc, kpn::in<"frame">, kpn::out<"scene">, "face_detector", 0> detector (detector_fn, 64);
|
||||||
|
kpn::ObjectNode<FaceAlignerFunc, kpn::in<"scene">, kpn::out<"aligned">, "face_aligner", 0> aligner (aligner_fn, 64);
|
||||||
|
kpn::ObjectNode<EmbedderFunc, kpn::in<"aligned">, kpn::out<"embedded">, "embedder", 0> embedder (embedder_fn, 32);
|
||||||
|
kpn::ObjectNode<FaceTrackerFunc, kpn::in<"embedded">, kpn::out<"tracked">, "face_tracker", 0> ftracker (ftracker_fn, 16);
|
||||||
|
kpn::ObjectNode<IdentityMatcherFunc, kpn::in<"tracked">, kpn::out<"matched">, "identity_matcher", 0> matcher (matcher_fn, 16);
|
||||||
|
kpn::ObjectNode<SceneTrackerFunc, kpn::in<"matched">, kpn::out<"annotation">, "scene_tracker", 0> tracker (tracker_fn, 16);
|
||||||
|
kpn::ObjectNode<ResultSinkFunc, kpn::in<"annotation">,kpn::out<>, "result_sink", 0> sink (sink_fn, 16);
|
||||||
|
|
||||||
|
// ── Build static network ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#ifdef SAE_DEBUG
|
||||||
|
kpn::ObjectNode<DebugRendererFunc, kpn::in<"matched">, kpn::out<>, "debug_renderer", 1> debug_node(debug_fn, 16);
|
||||||
|
|
||||||
|
// matcher → FanoutNode<MatchedSceneFrame,2> → scene_tracker + debug_node (auto-inserted)
|
||||||
|
auto net = kpn::make_network(
|
||||||
|
kpn::edge(source.output<"frame">(), detector.input<"frame">()),
|
||||||
|
kpn::edge(detector.output<"scene">(), aligner.input<"scene">()),
|
||||||
|
kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()),
|
||||||
|
kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()),
|
||||||
|
kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()),
|
||||||
|
kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()),
|
||||||
|
kpn::edge(matcher.output<"matched">(), debug_node.input<"matched">()),
|
||||||
|
kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">())
|
||||||
|
);
|
||||||
|
#else
|
||||||
|
auto net = kpn::make_network(
|
||||||
|
kpn::edge(source.output<"frame">(), detector.input<"frame">()),
|
||||||
|
kpn::edge(detector.output<"scene">(), aligner.input<"scene">()),
|
||||||
|
kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()),
|
||||||
|
kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()),
|
||||||
|
kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()),
|
||||||
|
kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()),
|
||||||
|
kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">())
|
||||||
|
);
|
||||||
|
#endif
|
||||||
|
|
||||||
|
// ── Run ───────────────────────────────────────────────────────────────────
|
||||||
|
// Note: StaticNetwork does not expose set_error_handler; node exceptions
|
||||||
|
// are printed to stderr by the KPN run_loop and the network continues.
|
||||||
|
std::cerr << "[main] starting pipeline…\n";
|
||||||
|
net.start();
|
||||||
|
|
||||||
|
// Main thread waits until ResultSinkFunc signals EOF completion
|
||||||
|
while (!done.load(std::memory_order_acquire))
|
||||||
|
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||||
|
|
||||||
|
net.stop();
|
||||||
|
net.print_diagnostics();
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
@@ -0,0 +1,150 @@
|
|||||||
|
#pragma once
|
||||||
|
#ifdef SAE_DEBUG
|
||||||
|
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <opencv2/imgcodecs.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
#include <filesystem>
|
||||||
|
#include <iostream>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
namespace fs = std::filesystem;
|
||||||
|
|
||||||
|
// ── DebugRendererFunc ─────────────────────────────────────────────────────────
|
||||||
|
// KPN sink node (SAE_DEBUG only): saves one directory of debug images per frame.
|
||||||
|
//
|
||||||
|
// Output layout:
|
||||||
|
// debug_frames/
|
||||||
|
// t0001.000/
|
||||||
|
// annotated.jpg original frame + coloured bboxes + name overlays
|
||||||
|
// brad_pitt_0.82.jpg 112×112 aligned crop | 1.5× context crop (side-by-side)
|
||||||
|
// unknown_0_0.77.jpg same for unidentified faces
|
||||||
|
//
|
||||||
|
// The node taps MatchedSceneFrame (before scene tracking) so every raw
|
||||||
|
// detection — including unknowns — is captured here.
|
||||||
|
|
||||||
|
struct DebugRendererFunc {
|
||||||
|
static constexpr std::string_view label() { return "debug_renderer"; }
|
||||||
|
|
||||||
|
explicit DebugRendererFunc(const Config& cfg)
|
||||||
|
: cfg_(cfg)
|
||||||
|
{
|
||||||
|
fs::create_directories(cfg_.debug_dir);
|
||||||
|
std::cerr << "[debug_renderer] output dir: " << cfg_.debug_dir << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
void operator()(MatchedSceneFrame mf) {
|
||||||
|
if (mf.source.eof || mf.source.image.empty()) return;
|
||||||
|
|
||||||
|
// Directory for this timestamp, e.g. "debug_frames/t0042.000/"
|
||||||
|
char buf[32];
|
||||||
|
std::snprintf(buf, sizeof(buf), "t%08.3f", mf.source.timestamp_sec);
|
||||||
|
fs::path dir = fs::path(cfg_.debug_dir) / buf;
|
||||||
|
fs::create_directories(dir);
|
||||||
|
|
||||||
|
// ── Annotated frame ───────────────────────────────────────────────────
|
||||||
|
cv::Mat annotated = mf.source.image.clone();
|
||||||
|
int unknown_idx = 0;
|
||||||
|
for (const auto& ia : mf.actors) {
|
||||||
|
bool known = (ia.actor_idx >= 0);
|
||||||
|
cv::Scalar colour = known
|
||||||
|
? cv::Scalar(0, 200, 60) // green for identified
|
||||||
|
: cv::Scalar(0, 100, 220); // orange for unknown
|
||||||
|
|
||||||
|
cv::Rect2f b = ia.bbox;
|
||||||
|
cv::rectangle(annotated, b, colour, 2);
|
||||||
|
|
||||||
|
std::string lbl = known
|
||||||
|
? (ia.name + " " + fmt_pct(ia.similarity))
|
||||||
|
: ("unknown " + fmt_pct(ia.similarity));
|
||||||
|
|
||||||
|
// Background strip for readability
|
||||||
|
int baseline = 0;
|
||||||
|
cv::Size ts = cv::getTextSize(lbl, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);
|
||||||
|
cv::Rect strip(static_cast<int>(b.x), static_cast<int>(b.y) - ts.height - 4,
|
||||||
|
ts.width + 4, ts.height + 6);
|
||||||
|
strip &= cv::Rect(0, 0, annotated.cols, annotated.rows);
|
||||||
|
if (strip.area() > 0)
|
||||||
|
cv::rectangle(annotated, strip, colour, cv::FILLED);
|
||||||
|
|
||||||
|
cv::putText(annotated, lbl,
|
||||||
|
cv::Point(static_cast<int>(b.x) + 2, static_cast<int>(b.y) - 2),
|
||||||
|
cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1,
|
||||||
|
cv::LINE_AA);
|
||||||
|
|
||||||
|
// ── Per-face crop image ───────────────────────────────────────────
|
||||||
|
// Left panel: 112×112 aligned crop. Right panel: expanded context.
|
||||||
|
cv::Mat panel = make_face_panel(mf.source.image, ia);
|
||||||
|
|
||||||
|
std::string stem = known
|
||||||
|
? (sanitise(ia.name) + "_" + fmt_sim(ia.similarity))
|
||||||
|
: ("unknown_" + std::to_string(unknown_idx++) + "_" + fmt_sim(ia.similarity));
|
||||||
|
cv::imwrite((dir / (stem + ".jpg")).string(), panel,
|
||||||
|
{cv::IMWRITE_JPEG_QUALITY, 90});
|
||||||
|
}
|
||||||
|
|
||||||
|
cv::imwrite((dir / "annotated.jpg").string(), annotated,
|
||||||
|
{cv::IMWRITE_JPEG_QUALITY, 90});
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
const Config& cfg_;
|
||||||
|
|
||||||
|
// Build a side-by-side panel: [112×112 aligned crop | context crop resized to 112×112]
|
||||||
|
cv::Mat make_face_panel(const cv::Mat& frame, const IdentifiedActor& ia) const {
|
||||||
|
// Left: aligned 112×112
|
||||||
|
cv::Mat left = ia.crop.empty()
|
||||||
|
? cv::Mat(112, 112, CV_8UC3, cv::Scalar(60, 60, 60))
|
||||||
|
: ia.crop.clone();
|
||||||
|
|
||||||
|
// Right: expanded bbox from original frame, resized to 112×112
|
||||||
|
cv::Rect2f expanded = expand_bbox(ia.bbox, cfg_.crop_context,
|
||||||
|
frame.cols, frame.rows);
|
||||||
|
cv::Mat right_raw = frame(expanded).clone();
|
||||||
|
cv::Mat right;
|
||||||
|
cv::resize(right_raw, right, {112, 112}, 0, 0, cv::INTER_LINEAR);
|
||||||
|
|
||||||
|
// Separator line
|
||||||
|
cv::Mat sep(112, 4, CV_8UC3, cv::Scalar(200, 200, 200));
|
||||||
|
|
||||||
|
cv::Mat panel;
|
||||||
|
cv::hconcat(std::vector<cv::Mat>{left, sep, right}, panel);
|
||||||
|
return panel;
|
||||||
|
}
|
||||||
|
|
||||||
|
static cv::Rect2f expand_bbox(cv::Rect2f b, float factor, int W, int H) {
|
||||||
|
float cx = b.x + b.width * 0.5f;
|
||||||
|
float cy = b.y + b.height * 0.5f;
|
||||||
|
float nw = b.width * factor;
|
||||||
|
float nh = b.height * factor;
|
||||||
|
float x = std::max(0.f, cx - nw * 0.5f);
|
||||||
|
float y = std::max(0.f, cy - nh * 0.5f);
|
||||||
|
nw = std::min(nw, (float)W - x);
|
||||||
|
nh = std::min(nh, (float)H - y);
|
||||||
|
return {x, y, nw, nh};
|
||||||
|
}
|
||||||
|
|
||||||
|
static std::string fmt_pct(float v) {
|
||||||
|
char buf[8];
|
||||||
|
std::snprintf(buf, sizeof(buf), "%.0f%%", v * 100.f);
|
||||||
|
return buf;
|
||||||
|
}
|
||||||
|
|
||||||
|
static std::string fmt_sim(float v) {
|
||||||
|
char buf[8];
|
||||||
|
std::snprintf(buf, sizeof(buf), "%.2f", v);
|
||||||
|
return buf;
|
||||||
|
}
|
||||||
|
|
||||||
|
static std::string sanitise(const std::string& s) {
|
||||||
|
std::string out;
|
||||||
|
out.reserve(s.size());
|
||||||
|
for (char c : s)
|
||||||
|
out += (std::isalnum(c) ? std::tolower(c) : '_');
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
#endif // SAE_DEBUG
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "arcface_embedder.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "trt_arcface_embedder.hpp"
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
// ── EmbedderFunc ──────────────────────────────────────────────────────────────
|
||||||
|
// KPN node: runs ArcFace on every 112×112 crop in an AlignedSceneFrame,
|
||||||
|
// producing one L2-normalised 512-dim embedding per face.
|
||||||
|
//
|
||||||
|
// Backend selection:
|
||||||
|
// --arcface-engine <path> → TrtArcFaceEmbedder (raw TensorRT, no ORT)
|
||||||
|
// otherwise → ArcFaceEmbedder (ONNX Runtime, picks best EP)
|
||||||
|
//
|
||||||
|
// All crops in one frame are batched into a single forward pass (capped at
|
||||||
|
// embed_batch_size). The backends serialise themselves; we only call them
|
||||||
|
// from the single embedder thread.
|
||||||
|
|
||||||
|
struct EmbedderFunc {
|
||||||
|
static constexpr std::string_view label() { return "embedder"; }
|
||||||
|
|
||||||
|
explicit EmbedderFunc(const Config& cfg, OrtProvider provider)
|
||||||
|
: batch_size_(std::max(1, cfg.embed_batch_size))
|
||||||
|
{
|
||||||
|
if (!cfg.arcface_engine.empty()) {
|
||||||
|
trt_ = std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
|
||||||
|
if (trt_->max_batch() < static_cast<int>(batch_size_))
|
||||||
|
throw std::runtime_error(
|
||||||
|
"embed_batch_size " + std::to_string(batch_size_) +
|
||||||
|
" exceeds engine max_batch " + std::to_string(trt_->max_batch()) +
|
||||||
|
" — rebuild engine with EMBED_BATCH=" + std::to_string(batch_size_));
|
||||||
|
} else {
|
||||||
|
ort_ = std::make_unique<ArcFaceEmbedder>(
|
||||||
|
cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
EmbeddedSceneFrame operator()(AlignedSceneFrame af) {
|
||||||
|
if (af.source.eof || af.crops.empty())
|
||||||
|
return {std::move(af.source), {}, {}, {}};
|
||||||
|
|
||||||
|
const auto& crops = af.crops;
|
||||||
|
std::vector<Embedding> embeddings;
|
||||||
|
embeddings.reserve(crops.size());
|
||||||
|
for (size_t i = 0; i < crops.size(); i += batch_size_) {
|
||||||
|
const size_t end = std::min(i + batch_size_, crops.size());
|
||||||
|
std::vector<cv::Mat> chunk_crops(crops.begin() + i, crops.begin() + end);
|
||||||
|
auto chunk = trt_ ? trt_->embed(chunk_crops) : ort_->embed(chunk_crops);
|
||||||
|
embeddings.insert(embeddings.end(), chunk.begin(), chunk.end());
|
||||||
|
}
|
||||||
|
|
||||||
|
return {std::move(af.source),
|
||||||
|
std::move(af.faces),
|
||||||
|
std::move(af.crops),
|
||||||
|
std::move(embeddings)};
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::unique_ptr<ArcFaceEmbedder> ort_;
|
||||||
|
std::unique_ptr<TrtArcFaceEmbedder> trt_;
|
||||||
|
size_t batch_size_;
|
||||||
|
};
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "face_utils.hpp"
|
||||||
|
|
||||||
|
#include <iostream>
|
||||||
|
|
||||||
|
// ── FaceAlignerFunc ───────────────────────────────────────────────────────────
|
||||||
|
// KPN node: applies a 5-point similarity transform to each detected face,
|
||||||
|
// producing a 112×112 BGR crop suitable for ArcFace inference.
|
||||||
|
//
|
||||||
|
// Alignment uses cv::estimateAffinePartial2D (RANSAC) to fit the detected
|
||||||
|
// landmarks to ArcFace canonical positions. Degenerate detections (where the
|
||||||
|
// affine fit fails) are silently dropped from the output vectors.
|
||||||
|
|
||||||
|
struct FaceAlignerFunc {
|
||||||
|
static constexpr std::string_view label() { return "face_aligner"; }
|
||||||
|
|
||||||
|
AlignedSceneFrame operator()(SceneFrame sf) {
|
||||||
|
if (sf.source.eof || sf.faces.empty())
|
||||||
|
return {std::move(sf.source), {}, {}};
|
||||||
|
|
||||||
|
std::vector<DetectedFace> good_faces;
|
||||||
|
std::vector<cv::Mat> crops;
|
||||||
|
good_faces.reserve(sf.faces.size());
|
||||||
|
crops.reserve(sf.faces.size());
|
||||||
|
|
||||||
|
for (auto& face : sf.faces) {
|
||||||
|
cv::Mat crop = align_face(sf.source.image, face.landmarks);
|
||||||
|
if (crop.empty()) {
|
||||||
|
std::cerr << "[face_aligner] degenerate detection skipped\n";
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
good_faces.push_back(face);
|
||||||
|
crops.push_back(std::move(crop));
|
||||||
|
}
|
||||||
|
|
||||||
|
return {std::move(sf.source), std::move(good_faces), std::move(crops)};
|
||||||
|
}
|
||||||
|
|
||||||
|
};
|
||||||
@@ -0,0 +1,61 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "scrfd_decoder.hpp"
|
||||||
|
#include "trt_scrfd_decoder.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
// ── FaceDetectorFunc ──────────────────────────────────────────────────────────
|
||||||
|
// KPN node: runs SCRFD-500MF to detect ALL faces in a frame.
|
||||||
|
//
|
||||||
|
// Backend selection:
|
||||||
|
// --detector-engine <path> → TrtScrfdDecoder (raw TensorRT, no ORT)
|
||||||
|
// otherwise → SCRFDDecoder (ONNX Runtime)
|
||||||
|
|
||||||
|
struct FaceDetectorFunc {
|
||||||
|
static constexpr std::string_view label() { return "face_detector"; }
|
||||||
|
|
||||||
|
explicit FaceDetectorFunc(const Config& cfg, OrtProvider provider)
|
||||||
|
: max_faces_(cfg.max_faces)
|
||||||
|
, min_face_px_(cfg.min_face_px)
|
||||||
|
{
|
||||||
|
if (!cfg.detector_engine.empty()) {
|
||||||
|
trt_ = std::make_unique<TrtScrfdDecoder>(
|
||||||
|
cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
|
||||||
|
} else {
|
||||||
|
ort_ = std::make_unique<SCRFDDecoder>(
|
||||||
|
cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
SceneFrame operator()(Frame f) {
|
||||||
|
if (f.eof) return {std::move(f), {}};
|
||||||
|
|
||||||
|
auto faces = trt_ ? trt_->detect(f.image) : ort_->detect(f.image);
|
||||||
|
|
||||||
|
// Drop faces below minimum pixel size (too small for reliable ArcFace alignment)
|
||||||
|
faces.erase(
|
||||||
|
std::remove_if(faces.begin(), faces.end(), [&](const DetectedFace& d) {
|
||||||
|
return d.bbox.width < min_face_px_ || d.bbox.height < min_face_px_;
|
||||||
|
}),
|
||||||
|
faces.end());
|
||||||
|
|
||||||
|
// Sort largest-first so max_faces_ keeps the most informative detections
|
||||||
|
std::sort(faces.begin(), faces.end(),
|
||||||
|
[](const DetectedFace& a, const DetectedFace& b) {
|
||||||
|
return a.bbox.area() > b.bbox.area();
|
||||||
|
});
|
||||||
|
if (static_cast<int>(faces.size()) > max_faces_)
|
||||||
|
faces.resize(max_faces_);
|
||||||
|
|
||||||
|
return {std::move(f), std::move(faces)};
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::unique_ptr<SCRFDDecoder> ort_;
|
||||||
|
std::unique_ptr<TrtScrfdDecoder> trt_;
|
||||||
|
int max_faces_{10};
|
||||||
|
float min_face_px_{40.f};
|
||||||
|
};
|
||||||
@@ -0,0 +1,243 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <algorithm>
|
||||||
|
#include <cmath>
|
||||||
|
#include <iostream>
|
||||||
|
#include <limits>
|
||||||
|
#include <map>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// ── FaceTrackerFunc ───────────────────────────────────────────────────────────
|
||||||
|
// KPN node: links face detections across consecutive frames using the Hungarian
|
||||||
|
// algorithm on a combined spatial (IoU) + embedding (cosine distance) cost.
|
||||||
|
//
|
||||||
|
// Each track accumulates a running directional mean of its ArcFace embeddings
|
||||||
|
// (averaged then re-normalised to the unit sphere). Once a track reaches
|
||||||
|
// min_frames observations its mean embedding is forwarded as track_embeddings[i]
|
||||||
|
// and track_mature[i] is set, allowing the identity matcher to use a cleaner,
|
||||||
|
// multi-frame signal instead of the noisy single-frame embedding.
|
||||||
|
//
|
||||||
|
// Assignment cost (track i, detection j):
|
||||||
|
// cost = alpha * (1 - IoU) + (1-alpha) * min(cosine_dist/2, 1)
|
||||||
|
// Gated to INF when IoU < min_iou AND cosine_dist > max_embed_dist.
|
||||||
|
//
|
||||||
|
// Unmatched tracks have their frames_missing counter incremented; they are
|
||||||
|
// expired once frames_missing > max_frames_missing.
|
||||||
|
|
||||||
|
struct FaceTrackerFunc {
|
||||||
|
static constexpr std::string_view label() { return "face_tracker"; }
|
||||||
|
|
||||||
|
struct TrackState {
|
||||||
|
cv::Rect2f bbox;
|
||||||
|
Embedding mean_emb{};
|
||||||
|
int n_frames{0};
|
||||||
|
int frames_missing{0};
|
||||||
|
};
|
||||||
|
|
||||||
|
explicit FaceTrackerFunc(const Config& cfg)
|
||||||
|
: alpha_(cfg.track_alpha)
|
||||||
|
, min_iou_(cfg.track_min_iou)
|
||||||
|
, max_embed_dist_(cfg.track_max_embed_dist)
|
||||||
|
, max_missing_(cfg.track_max_frames_missing)
|
||||||
|
, min_frames_(cfg.track_min_frames)
|
||||||
|
{
|
||||||
|
std::cerr << "[face_tracker] alpha=" << alpha_
|
||||||
|
<< " min_iou=" << min_iou_
|
||||||
|
<< " max_embed_dist=" << max_embed_dist_
|
||||||
|
<< " max_missing=" << max_missing_
|
||||||
|
<< " min_frames=" << min_frames_ << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
TrackedSceneFrame operator()(EmbeddedSceneFrame ef) {
|
||||||
|
if (ef.source.eof) {
|
||||||
|
tracks_.clear();
|
||||||
|
TrackedSceneFrame out;
|
||||||
|
out.source = std::move(ef.source);
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
|
||||||
|
const int n_det = static_cast<int>(ef.embeddings.size());
|
||||||
|
|
||||||
|
if (ef.source.is_cut && !tracks_.empty()) {
|
||||||
|
std::cerr << "[face_tracker] cut — clearing " << tracks_.size() << " tracks\n";
|
||||||
|
tracks_.clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Snapshot active track IDs so the map can be modified safely below
|
||||||
|
std::vector<int> tids;
|
||||||
|
tids.reserve(tracks_.size());
|
||||||
|
for (auto& [tid, _] : tracks_) tids.push_back(tid);
|
||||||
|
const int n_trk = static_cast<int>(tids.size());
|
||||||
|
|
||||||
|
// ── Cost matrix [n_trk × n_det] ──────────────────────────────────────
|
||||||
|
constexpr float INF_COST = 1e6f;
|
||||||
|
std::vector<std::vector<float>> cost(n_trk,
|
||||||
|
std::vector<float>(n_det, INF_COST));
|
||||||
|
|
||||||
|
for (int ti = 0; ti < n_trk; ++ti) {
|
||||||
|
const TrackState& ts = tracks_[tids[ti]];
|
||||||
|
for (int di = 0; di < n_det; ++di) {
|
||||||
|
float iou_v = iou(ts.bbox, ef.faces[di].bbox);
|
||||||
|
float emb_d = (ts.n_frames > 0)
|
||||||
|
? 1.f - cosine_similarity(ts.mean_emb, ef.embeddings[di])
|
||||||
|
: 1.f;
|
||||||
|
if (iou_v < min_iou_ && emb_d > max_embed_dist_) continue;
|
||||||
|
float s = 1.f - iou_v;
|
||||||
|
float e = std::min(emb_d * 0.5f, 1.f);
|
||||||
|
cost[ti][di] = alpha_ * s + (1.f - alpha_) * e;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Hungarian assignment ──────────────────────────────────────────────
|
||||||
|
std::vector<int> assign(n_trk, -1);
|
||||||
|
if (n_trk > 0 && n_det > 0)
|
||||||
|
assign = hungarian(cost, n_trk, n_det);
|
||||||
|
|
||||||
|
// ── Build output frame ────────────────────────────────────────────────
|
||||||
|
TrackedSceneFrame out;
|
||||||
|
out.source = ef.source;
|
||||||
|
out.faces = ef.faces;
|
||||||
|
out.crops = ef.crops;
|
||||||
|
out.embeddings = ef.embeddings;
|
||||||
|
out.track_ids.assign(n_det, -1);
|
||||||
|
out.track_embeddings = ef.embeddings; // default: per-frame embedding
|
||||||
|
out.track_mature.assign(n_det, false);
|
||||||
|
|
||||||
|
std::vector<bool> det_matched(n_det, false);
|
||||||
|
|
||||||
|
// Update matched tracks
|
||||||
|
for (int ti = 0; ti < n_trk; ++ti) {
|
||||||
|
int di = assign[ti];
|
||||||
|
bool valid = (di >= 0 && di < n_det && cost[ti][di] < INF_COST * 0.5f);
|
||||||
|
TrackState& ts = tracks_[tids[ti]];
|
||||||
|
if (!valid) {
|
||||||
|
ts.frames_missing++;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
update_mean(ts.mean_emb, ts.n_frames, ef.embeddings[di]);
|
||||||
|
ts.bbox = ef.faces[di].bbox;
|
||||||
|
ts.n_frames++;
|
||||||
|
ts.frames_missing = 0;
|
||||||
|
det_matched[di] = true;
|
||||||
|
|
||||||
|
out.track_ids[di] = tids[ti];
|
||||||
|
out.track_embeddings[di] = ts.mean_emb;
|
||||||
|
out.track_mature[di] = (ts.n_frames >= min_frames_);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create new tracks for unmatched detections
|
||||||
|
for (int di = 0; di < n_det; ++di) {
|
||||||
|
if (det_matched[di]) continue;
|
||||||
|
int tid = next_id_++;
|
||||||
|
TrackState ts;
|
||||||
|
ts.bbox = ef.faces[di].bbox;
|
||||||
|
ts.mean_emb = ef.embeddings[di];
|
||||||
|
ts.n_frames = 1;
|
||||||
|
tracks_[tid] = ts;
|
||||||
|
out.track_ids[di] = tid;
|
||||||
|
// track_embeddings[di] already initialised to per-frame embedding
|
||||||
|
}
|
||||||
|
|
||||||
|
// Expire stale tracks
|
||||||
|
for (auto it = tracks_.begin(); it != tracks_.end(); ) {
|
||||||
|
it = (it->second.frames_missing > max_missing_)
|
||||||
|
? tracks_.erase(it) : std::next(it);
|
||||||
|
}
|
||||||
|
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
// IoU of two axis-aligned bounding boxes
|
||||||
|
static float iou(const cv::Rect2f& a, const cv::Rect2f& b) {
|
||||||
|
float ix = std::max(0.f, std::min(a.x + a.width, b.x + b.width)
|
||||||
|
- std::max(a.x, b.x));
|
||||||
|
float iy = std::max(0.f, std::min(a.y + a.height, b.y + b.height)
|
||||||
|
- std::max(a.y, b.y));
|
||||||
|
float inter = ix * iy;
|
||||||
|
if (inter <= 0.f) return 0.f;
|
||||||
|
return inter / (a.width * a.height + b.width * b.height - inter);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Online directional mean: average then re-normalise to unit sphere
|
||||||
|
static void update_mean(Embedding& mean, int n_prev, const Embedding& emb) {
|
||||||
|
float norm_sq = 0.f;
|
||||||
|
for (int k = 0; k < 512; ++k) {
|
||||||
|
mean[k] = (mean[k] * n_prev + emb[k]) / (n_prev + 1);
|
||||||
|
norm_sq += mean[k] * mean[k];
|
||||||
|
}
|
||||||
|
float inv = 1.f / std::sqrt(norm_sq);
|
||||||
|
for (int k = 0; k < 512; ++k) mean[k] *= inv;
|
||||||
|
}
|
||||||
|
|
||||||
|
// O(n³) potential-based Hungarian algorithm (Jonker-Volgenant / Kuhn-Munkres).
|
||||||
|
// Returns assign[row] = col (0-indexed), or -1 when row is matched to a
|
||||||
|
// padded virtual column (i.e., unmatched). Rectangular matrices are padded
|
||||||
|
// to square with 0-cost virtual entries so leftover rows/cols are absorbed
|
||||||
|
// cheaply rather than being forced onto real rows/cols.
|
||||||
|
static std::vector<int> hungarian(
|
||||||
|
const std::vector<std::vector<float>>& C, int nr, int nc)
|
||||||
|
{
|
||||||
|
const int N = std::max(nr, nc);
|
||||||
|
constexpr float INF_VAL = 1e30f;
|
||||||
|
|
||||||
|
// Expand to N×N, filling virtual entries with 0
|
||||||
|
std::vector<std::vector<float>> sq(N, std::vector<float>(N, 0.f));
|
||||||
|
for (int i = 0; i < nr; ++i)
|
||||||
|
for (int j = 0; j < nc; ++j)
|
||||||
|
sq[i][j] = C[i][j];
|
||||||
|
|
||||||
|
std::vector<float> u(N + 1, 0.f), v(N + 1, 0.f);
|
||||||
|
std::vector<int> p(N + 1, 0), way(N + 1, 0);
|
||||||
|
|
||||||
|
for (int i = 1; i <= N; ++i) {
|
||||||
|
p[0] = i;
|
||||||
|
int j0 = 0;
|
||||||
|
std::vector<float> minv(N + 1, INF_VAL);
|
||||||
|
std::vector<bool> used(N + 1, false);
|
||||||
|
do {
|
||||||
|
used[j0] = true;
|
||||||
|
int i0 = p[j0], j1 = -1;
|
||||||
|
float delta = INF_VAL;
|
||||||
|
for (int j = 1; j <= N; ++j) {
|
||||||
|
if (!used[j]) {
|
||||||
|
float cur = sq[i0-1][j-1] - u[i0] - v[j];
|
||||||
|
if (cur < minv[j]) { minv[j] = cur; way[j] = j0; }
|
||||||
|
if (minv[j] < delta) { delta = minv[j]; j1 = j; }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for (int j = 0; j <= N; ++j) {
|
||||||
|
if (used[j]) { u[p[j]] += delta; v[j] -= delta; }
|
||||||
|
else minv[j] -= delta;
|
||||||
|
}
|
||||||
|
j0 = j1;
|
||||||
|
} while (p[j0] != 0);
|
||||||
|
do {
|
||||||
|
int j1 = way[j0];
|
||||||
|
p[j0] = p[j1];
|
||||||
|
j0 = j1;
|
||||||
|
} while (j0);
|
||||||
|
}
|
||||||
|
|
||||||
|
// p[j] = row (1-indexed) assigned to column j (1-indexed)
|
||||||
|
std::vector<int> ans(nr, -1);
|
||||||
|
for (int j = 1; j <= N; ++j) {
|
||||||
|
int row = p[j] - 1;
|
||||||
|
int col = j - 1;
|
||||||
|
if (row >= 0 && row < nr && col < nc)
|
||||||
|
ans[row] = col;
|
||||||
|
// col >= nc → virtual column → row stays unmatched (-1)
|
||||||
|
}
|
||||||
|
return ans;
|
||||||
|
}
|
||||||
|
|
||||||
|
std::map<int, TrackState> tracks_;
|
||||||
|
int next_id_{0};
|
||||||
|
float alpha_;
|
||||||
|
float min_iou_;
|
||||||
|
float max_embed_dist_;
|
||||||
|
int max_missing_;
|
||||||
|
int min_frames_;
|
||||||
|
};
|
||||||
@@ -0,0 +1,144 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
#include "ffmpeg_decoder.hpp"
|
||||||
|
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
#include <chrono>
|
||||||
|
#include <iostream>
|
||||||
|
#include <memory>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <thread>
|
||||||
|
|
||||||
|
// ── FrameSourceFunc ───────────────────────────────────────────────────────────
|
||||||
|
// KPN source node: reads a movie file and emits one Frame per sample interval.
|
||||||
|
//
|
||||||
|
// Decode backend: FFmpeg with NVDEC (_cuvid) when available, CPU otherwise.
|
||||||
|
//
|
||||||
|
// Sampling strategy: seek to the next target timestamp rather than decoding
|
||||||
|
// every frame, which is fast even for 1-FPS sampling of a 2-hour film.
|
||||||
|
//
|
||||||
|
// EOF handling: when the movie ends, emits a Frame with eof=true, then sleeps
|
||||||
|
// 500 ms between subsequent calls until the KPN network stops the thread.
|
||||||
|
|
||||||
|
struct FrameSourceFunc {
|
||||||
|
static constexpr std::string_view label() { return "frame_source"; }
|
||||||
|
|
||||||
|
explicit FrameSourceFunc(const Config& cfg)
|
||||||
|
: decoder_(std::make_unique<FFmpegDecoder>(cfg.movie_path))
|
||||||
|
{
|
||||||
|
sample_interval_sec_ = 1.0 / cfg.sample_fps;
|
||||||
|
next_pos_sec_ = cfg.start_sec;
|
||||||
|
end_sec_ = cfg.end_sec;
|
||||||
|
cut_threshold_ = cfg.cut_threshold;
|
||||||
|
max_decode_fps_ = cfg.max_decode_fps;
|
||||||
|
|
||||||
|
double total_s = decoder_->duration_sec();
|
||||||
|
double span_s = (end_sec_ > 0 ? std::min(end_sec_, total_s) : total_s)
|
||||||
|
- cfg.start_sec;
|
||||||
|
int n_frames = static_cast<int>(span_s * cfg.sample_fps);
|
||||||
|
std::cerr << "[frame_source] decoder=" << decoder_->codec_name()
|
||||||
|
<< " (" << (decoder_->hw_active() ? "NVDEC" : "CPU") << ")"
|
||||||
|
<< " video_fps=" << decoder_->fps()
|
||||||
|
<< " start=" << cfg.start_sec << "s"
|
||||||
|
<< (end_sec_ > 0 ? " end=" + std::to_string(end_sec_) + "s" : "")
|
||||||
|
<< " sample_fps=" << cfg.sample_fps
|
||||||
|
<< " frames_to_emit=" << n_frames << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
Frame operator()() {
|
||||||
|
if (hit_eof_) {
|
||||||
|
std::this_thread::sleep_for(std::chrono::milliseconds(500));
|
||||||
|
return Frame{{}, 0.0, -1, /*eof=*/true};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wall-clock rate cap. KPN source nodes resubmit immediately on push
|
||||||
|
// overflow, with no backpressure; without this cap we'd decode-and-drop
|
||||||
|
// in a tight loop whenever downstream stalls. The cap also protects
|
||||||
|
// ORT-only deployments where the pipeline can't keep up at decode speed.
|
||||||
|
if (max_decode_fps_ > 0.f) {
|
||||||
|
const auto now = std::chrono::steady_clock::now();
|
||||||
|
if (!rate_started_) {
|
||||||
|
rate_started_ = true;
|
||||||
|
next_decode_at_ = now;
|
||||||
|
}
|
||||||
|
if (now < next_decode_at_)
|
||||||
|
std::this_thread::sleep_until(next_decode_at_);
|
||||||
|
const auto period = std::chrono::nanoseconds(
|
||||||
|
static_cast<int64_t>(1e9f / max_decode_fps_));
|
||||||
|
// Anchor the next slot off the slot we just consumed, not off
|
||||||
|
// wall-clock now() — keeps the average rate stable. If we fell
|
||||||
|
// behind by more than one period, snap forward to avoid building
|
||||||
|
// up an unbounded sleep debt.
|
||||||
|
next_decode_at_ += period;
|
||||||
|
if (next_decode_at_ < now)
|
||||||
|
next_decode_at_ = now + period;
|
||||||
|
}
|
||||||
|
|
||||||
|
auto t0 = std::chrono::steady_clock::now();
|
||||||
|
cv::Mat img = decoder_->read_at(next_pos_sec_);
|
||||||
|
auto t1 = std::chrono::steady_clock::now();
|
||||||
|
double decode_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
|
||||||
|
decode_ms_acc_ += decode_ms;
|
||||||
|
++decode_count_;
|
||||||
|
if (decode_count_ % 10 == 0) {
|
||||||
|
double avg_ms = decode_ms_acc_ / 10.0;
|
||||||
|
double avg_fps = avg_ms > 0.0 ? 1000.0 / avg_ms : 0.0;
|
||||||
|
std::cerr << "[frame_source] decode avg=" << avg_ms << "ms"
|
||||||
|
<< " fps=" << avg_fps << "\n";
|
||||||
|
decode_ms_acc_ = 0.0;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (img.empty()) {
|
||||||
|
hit_eof_ = true;
|
||||||
|
std::cerr << "[frame_source] EOF at t=" << next_pos_sec_ << "s\n";
|
||||||
|
return Frame{{}, next_pos_sec_, frame_idx_++, /*eof=*/true};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cut detection: compare grayscale histogram to previous frame
|
||||||
|
bool is_cut = false;
|
||||||
|
cv::Mat gray;
|
||||||
|
cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);
|
||||||
|
cv::Mat hist;
|
||||||
|
const int bins = 64;
|
||||||
|
const float range[] = {0.f, 256.f};
|
||||||
|
const float* ranges = range;
|
||||||
|
cv::calcHist(&gray, 1, nullptr, cv::Mat(), hist, 1, &bins, &ranges);
|
||||||
|
cv::normalize(hist, hist, 1.0, 0.0, cv::NORM_L1);
|
||||||
|
|
||||||
|
if (prev_hist_valid_) {
|
||||||
|
double corr = cv::compareHist(prev_hist_, hist, cv::HISTCMP_CORREL);
|
||||||
|
is_cut = (corr < cut_threshold_);
|
||||||
|
if (is_cut)
|
||||||
|
std::cerr << "[frame_source] cut at t=" << next_pos_sec_
|
||||||
|
<< "s hist_corr=" << corr << "\n";
|
||||||
|
}
|
||||||
|
prev_hist_ = hist;
|
||||||
|
prev_hist_valid_ = true;
|
||||||
|
|
||||||
|
Frame f{img, next_pos_sec_, frame_idx_++, /*eof=*/false, is_cut};
|
||||||
|
next_pos_sec_ += sample_interval_sec_;
|
||||||
|
if (end_sec_ > 0 && next_pos_sec_ > end_sec_) {
|
||||||
|
hit_eof_ = true;
|
||||||
|
std::cerr << "[frame_source] reached end_sec=" << end_sec_ << "s\n";
|
||||||
|
}
|
||||||
|
return f;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::unique_ptr<FFmpegDecoder> decoder_;
|
||||||
|
double sample_interval_sec_{1.0};
|
||||||
|
double next_pos_sec_{0.0};
|
||||||
|
double end_sec_{-1.0};
|
||||||
|
float cut_threshold_{0.70f};
|
||||||
|
float max_decode_fps_{0.f};
|
||||||
|
std::chrono::steady_clock::time_point next_decode_at_{};
|
||||||
|
bool rate_started_{false};
|
||||||
|
int64_t frame_idx_{0};
|
||||||
|
bool hit_eof_{false};
|
||||||
|
cv::Mat prev_hist_;
|
||||||
|
bool prev_hist_valid_{false};
|
||||||
|
double decode_ms_acc_{0.0};
|
||||||
|
int decode_count_{0};
|
||||||
|
};
|
||||||
@@ -0,0 +1,154 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
#include "gallery/gallery_store.hpp"
|
||||||
|
#include "gallery/gallery_calibration.hpp"
|
||||||
|
|
||||||
|
#include <cmath>
|
||||||
|
#include <limits>
|
||||||
|
#include <iostream>
|
||||||
|
|
||||||
|
// ── IdentityMatcherFunc ───────────────────────────────────────────────────────
|
||||||
|
// KPN node: compares each embedding against every reference embedding in the
|
||||||
|
// actor gallery using cosine similarity.
|
||||||
|
//
|
||||||
|
// Matching strategy — two modes selected at construction time:
|
||||||
|
//
|
||||||
|
// Calibrated (preferred): gallery calibration fits a sigmoid
|
||||||
|
// P(match) = σ(a·similarity + b) from intra/inter-class pairs.
|
||||||
|
// A face is accepted if P(match | best_actor) > prob_threshold.
|
||||||
|
//
|
||||||
|
// Fallback (no calibration): dual-criterion accept —
|
||||||
|
// (a) best cosine distance < match_threshold, OR
|
||||||
|
// (b) ratio test: best_dist/second_best_dist < match_ratio
|
||||||
|
// AND best_dist < match_ratio_ceil.
|
||||||
|
//
|
||||||
|
// In both modes, per-actor best similarity is determined by scanning all
|
||||||
|
// reference embeddings and taking the closest (best-of-N).
|
||||||
|
|
||||||
|
struct IdentityMatcherFunc {
|
||||||
|
static constexpr std::string_view label() { return "identity_matcher"; }
|
||||||
|
|
||||||
|
IdentityMatcherFunc(const ActorGallery& gallery, const Config& cfg)
|
||||||
|
: gallery_(gallery)
|
||||||
|
, prob_threshold_(cfg.prob_threshold)
|
||||||
|
, log_prior_odds_(std::log(cfg.match_prior / (1.f - cfg.match_prior)))
|
||||||
|
, threshold_(cfg.match_threshold)
|
||||||
|
, ratio_(cfg.match_ratio)
|
||||||
|
, ratio_ceil_(cfg.match_ratio_ceil)
|
||||||
|
{
|
||||||
|
for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
|
||||||
|
for (const auto& emb : gallery_.actors[ai].embeddings) {
|
||||||
|
flat_emb_.push_back(emb);
|
||||||
|
flat_actor_.push_back(ai);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cal_ = calibrate_gallery(flat_emb_, flat_actor_);
|
||||||
|
|
||||||
|
if (cal_.valid) {
|
||||||
|
std::cerr << "[identity_matcher] calibrated Bayesian matching"
|
||||||
|
<< " prior=" << cfg.match_prior
|
||||||
|
<< " P_threshold=" << prob_threshold_
|
||||||
|
<< " effective_sim_boundary="
|
||||||
|
<< cal_.boundary_at(prob_threshold_, log_prior_odds_) << "\n";
|
||||||
|
} else {
|
||||||
|
std::cerr << "[identity_matcher] threshold matching (calibration skipped)"
|
||||||
|
<< " threshold=" << threshold_
|
||||||
|
<< " ratio=" << ratio_ << " ratio_ceil=" << ratio_ceil_ << "\n";
|
||||||
|
}
|
||||||
|
std::cerr << "[identity_matcher] gallery: "
|
||||||
|
<< gallery_.actors.size() << " actors, "
|
||||||
|
<< flat_emb_.size() << " reference embeddings\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
MatchedSceneFrame operator()(TrackedSceneFrame tf) {
|
||||||
|
if (tf.source.eof) return {std::move(tf.source), {}};
|
||||||
|
|
||||||
|
std::vector<IdentifiedActor> actors;
|
||||||
|
actors.reserve(tf.embeddings.size());
|
||||||
|
|
||||||
|
for (int fi = 0; fi < static_cast<int>(tf.embeddings.size()); ++fi) {
|
||||||
|
// Prefer the track's accumulated mean embedding when the track is
|
||||||
|
// mature (≥ min_frames observations) — more stable than single-frame.
|
||||||
|
const Embedding& query = tf.track_mature[fi]
|
||||||
|
? tf.track_embeddings[fi]
|
||||||
|
: tf.embeddings[fi];
|
||||||
|
|
||||||
|
// Per-actor best cosine similarity (max dot product)
|
||||||
|
std::vector<float> best_sim(gallery_.actors.size(),
|
||||||
|
-std::numeric_limits<float>::max());
|
||||||
|
|
||||||
|
for (int ei = 0; ei < static_cast<int>(flat_emb_.size()); ++ei) {
|
||||||
|
float sim = cosine_similarity(query, flat_emb_[ei]);
|
||||||
|
int ai = flat_actor_[ei];
|
||||||
|
if (sim > best_sim[ai]) best_sim[ai] = sim;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find best and second-best actor by similarity
|
||||||
|
int best_actor = -1;
|
||||||
|
int second_actor = -1;
|
||||||
|
float best_s = -std::numeric_limits<float>::max();
|
||||||
|
float second_s = -std::numeric_limits<float>::max();
|
||||||
|
for (int ai = 0; ai < static_cast<int>(best_sim.size()); ++ai) {
|
||||||
|
if (best_sim[ai] > best_s) {
|
||||||
|
second_s = best_s;
|
||||||
|
second_actor = best_actor;
|
||||||
|
best_s = best_sim[ai];
|
||||||
|
best_actor = ai;
|
||||||
|
} else if (best_sim[ai] > second_s) {
|
||||||
|
second_s = best_sim[ai];
|
||||||
|
second_actor = ai;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
(void)second_actor;
|
||||||
|
|
||||||
|
bool accept = false;
|
||||||
|
if (best_actor >= 0) {
|
||||||
|
if (cal_.valid) {
|
||||||
|
accept = cal_.probability(best_s, log_prior_odds_) > prob_threshold_;
|
||||||
|
} else {
|
||||||
|
float best_d = 1.f - best_s;
|
||||||
|
float second_d = (second_s > -std::numeric_limits<float>::max())
|
||||||
|
? 1.f - second_s
|
||||||
|
: std::numeric_limits<float>::max();
|
||||||
|
bool absolute = best_d < threshold_;
|
||||||
|
bool ratio = (best_d < ratio_ceil_) &&
|
||||||
|
(second_d == std::numeric_limits<float>::max() ||
|
||||||
|
best_d / second_d < ratio_);
|
||||||
|
accept = absolute || ratio;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
IdentifiedActor ia;
|
||||||
|
ia.bbox = tf.faces[fi].bbox;
|
||||||
|
ia.crop = tf.crops[fi];
|
||||||
|
ia.track_id = tf.track_ids[fi];
|
||||||
|
|
||||||
|
if (accept) {
|
||||||
|
ia.actor_idx = best_actor;
|
||||||
|
ia.name = gallery_.actors[best_actor].name;
|
||||||
|
ia.imdb_id = gallery_.actors[best_actor].imdb_id;
|
||||||
|
ia.similarity = cal_.valid
|
||||||
|
? cal_.probability(best_s, log_prior_odds_)
|
||||||
|
: best_s;
|
||||||
|
}
|
||||||
|
// actor_idx == -1, name == "" → unknown face
|
||||||
|
|
||||||
|
actors.push_back(std::move(ia));
|
||||||
|
}
|
||||||
|
|
||||||
|
return {std::move(tf.source), std::move(actors)};
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
ActorGallery gallery_;
|
||||||
|
GalleryCalibration cal_;
|
||||||
|
float prob_threshold_;
|
||||||
|
float log_prior_odds_;
|
||||||
|
float threshold_;
|
||||||
|
float ratio_;
|
||||||
|
float ratio_ceil_;
|
||||||
|
std::vector<Embedding> flat_emb_;
|
||||||
|
std::vector<int> flat_actor_;
|
||||||
|
};
|
||||||
@@ -0,0 +1,140 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <kpn/main_thread_node.hpp>
|
||||||
|
|
||||||
|
#include <opencv2/highgui.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
// ── PreviewNode ───────────────────────────────────────────────────────────────
|
||||||
|
// MainThreadNode: receives MatchedSceneFrame, draws annotations and shows the
|
||||||
|
// frame in an OpenCV window. Runs on the main thread via preview.step().
|
||||||
|
//
|
||||||
|
// Returns false on EOF (ends the main loop) or when 'q' / Escape is pressed.
|
||||||
|
|
||||||
|
class PreviewNode : public kpn::MainThreadNode<PreviewNode,
|
||||||
|
kpn::in<"matched">,
|
||||||
|
MatchedSceneFrame> {
|
||||||
|
public:
|
||||||
|
static constexpr std::string_view label() { return "preview"; }
|
||||||
|
explicit PreviewNode(const Config& cfg, std::size_t fifo_capacity = 4)
|
||||||
|
: kpn::MainThreadNode<PreviewNode, kpn::in<"matched">, MatchedSceneFrame>(fifo_capacity)
|
||||||
|
, max_display_w_(1280)
|
||||||
|
{
|
||||||
|
cv::namedWindow("scene_preview", cv::WINDOW_NORMAL);
|
||||||
|
cv::resizeWindow("scene_preview", 1280, 720);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Called by step() on the main thread for each ready MatchedSceneFrame.
|
||||||
|
bool operator()(MatchedSceneFrame mf) {
|
||||||
|
if (mf.source.eof) return false;
|
||||||
|
|
||||||
|
cv::Mat display = mf.source.image.clone();
|
||||||
|
draw_detections(display, mf.actors);
|
||||||
|
draw_hud(display, mf.source.timestamp_sec, mf.actors);
|
||||||
|
|
||||||
|
// Fit to display width while keeping aspect ratio
|
||||||
|
if (display.cols > max_display_w_) {
|
||||||
|
float scale = static_cast<float>(max_display_w_) / display.cols;
|
||||||
|
cv::resize(display, display, {}, scale, scale, cv::INTER_AREA);
|
||||||
|
}
|
||||||
|
|
||||||
|
cv::imshow("scene_preview", display);
|
||||||
|
|
||||||
|
int key = cv::waitKey(1) & 0xFF;
|
||||||
|
return (key != 'q' && key != 27 /* Esc */);
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
int max_display_w_;
|
||||||
|
|
||||||
|
static void draw_detections(cv::Mat& img,
|
||||||
|
const std::vector<IdentifiedActor>& actors) {
|
||||||
|
for (const auto& ia : actors) {
|
||||||
|
bool known = (ia.actor_idx >= 0);
|
||||||
|
|
||||||
|
// Green for identified, orange for unknown
|
||||||
|
cv::Scalar box_colour = known
|
||||||
|
? cv::Scalar(0, 210, 60)
|
||||||
|
: cv::Scalar(0, 140, 255);
|
||||||
|
|
||||||
|
// Scale bbox to display image
|
||||||
|
cv::Rect2f b = ia.bbox;
|
||||||
|
cv::rectangle(img, b, box_colour, 2, cv::LINE_AA);
|
||||||
|
|
||||||
|
std::string tid = (ia.track_id >= 0) ? (" #" + std::to_string(ia.track_id)) : "";
|
||||||
|
std::string label = known
|
||||||
|
? (ia.name + tid + " " + pct(ia.similarity))
|
||||||
|
: ("?" + tid);
|
||||||
|
|
||||||
|
// Dark backing strip so text is readable on any background
|
||||||
|
int baseline = 0;
|
||||||
|
cv::Size ts = cv::getTextSize(label, cv::FONT_HERSHEY_DUPLEX,
|
||||||
|
0.55, 1, &baseline);
|
||||||
|
cv::Point tl(static_cast<int>(b.x),
|
||||||
|
std::max(0, static_cast<int>(b.y) - ts.height - 6));
|
||||||
|
cv::Rect backing(tl.x, tl.y, ts.width + 8, ts.height + 8);
|
||||||
|
backing &= cv::Rect(0, 0, img.cols, img.rows);
|
||||||
|
if (backing.area() > 0)
|
||||||
|
cv::rectangle(img, backing, box_colour * 0.6, cv::FILLED);
|
||||||
|
|
||||||
|
cv::putText(img, label,
|
||||||
|
cv::Point(tl.x + 4, tl.y + ts.height + 2),
|
||||||
|
cv::FONT_HERSHEY_DUPLEX, 0.55,
|
||||||
|
cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
static void draw_hud(cv::Mat& img, double timestamp_sec,
|
||||||
|
const std::vector<IdentifiedActor>& actors) {
|
||||||
|
int known = 0;
|
||||||
|
int unknown = 0;
|
||||||
|
for (const auto& a : actors) (a.actor_idx >= 0 ? known : unknown)++;
|
||||||
|
|
||||||
|
// Timestamp and actor count in top-left corner
|
||||||
|
char buf[128];
|
||||||
|
std::snprintf(buf, sizeof(buf),
|
||||||
|
"t = %dm %02ds | %d identified %d unknown",
|
||||||
|
static_cast<int>(timestamp_sec) / 60,
|
||||||
|
static_cast<int>(timestamp_sec) % 60,
|
||||||
|
known, unknown);
|
||||||
|
|
||||||
|
int baseline = 0;
|
||||||
|
cv::Size ts = cv::getTextSize(buf, cv::FONT_HERSHEY_SIMPLEX,
|
||||||
|
0.6, 1, &baseline);
|
||||||
|
cv::Rect hud(0, 0, ts.width + 16, ts.height + 12);
|
||||||
|
hud &= cv::Rect(0, 0, img.cols, img.rows);
|
||||||
|
cv::rectangle(img, hud, cv::Scalar(20, 20, 20), cv::FILLED);
|
||||||
|
cv::putText(img, buf, cv::Point(8, ts.height + 6),
|
||||||
|
cv::FONT_HERSHEY_SIMPLEX, 0.6,
|
||||||
|
cv::Scalar(220, 220, 220), 1, cv::LINE_AA);
|
||||||
|
|
||||||
|
// Active actor name strip along the bottom
|
||||||
|
if (known > 0) {
|
||||||
|
std::string names;
|
||||||
|
for (const auto& a : actors) {
|
||||||
|
if (a.actor_idx < 0) continue;
|
||||||
|
if (!names.empty()) names += " ";
|
||||||
|
names += a.name;
|
||||||
|
}
|
||||||
|
cv::Size ns = cv::getTextSize(names, cv::FONT_HERSHEY_SIMPLEX,
|
||||||
|
0.55, 1, &baseline);
|
||||||
|
int y = img.rows - ns.height - 10;
|
||||||
|
cv::Rect strip(0, y - 6, img.cols, ns.height + 16);
|
||||||
|
strip &= cv::Rect(0, 0, img.cols, img.rows);
|
||||||
|
cv::rectangle(img, strip, cv::Scalar(10, 10, 10, 180), cv::FILLED);
|
||||||
|
cv::putText(img, names, cv::Point(8, img.rows - 10),
|
||||||
|
cv::FONT_HERSHEY_SIMPLEX, 0.55,
|
||||||
|
cv::Scalar(80, 220, 100), 1, cv::LINE_AA);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
static std::string pct(float v) {
|
||||||
|
char buf[8];
|
||||||
|
std::snprintf(buf, sizeof(buf), "%.0f%%", v * 100.f);
|
||||||
|
return buf;
|
||||||
|
}
|
||||||
|
};
|
||||||
@@ -0,0 +1,204 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <nlohmann/json.hpp>
|
||||||
|
#include <algorithm>
|
||||||
|
#include <atomic>
|
||||||
|
#include <cmath>
|
||||||
|
#include <fstream>
|
||||||
|
#include <iostream>
|
||||||
|
#include <map>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
using json = nlohmann::json;
|
||||||
|
|
||||||
|
// ── ResultSinkFunc ────────────────────────────────────────────────────────────
|
||||||
|
// KPN sink node: accumulates SceneAnnotations and writes the final JSON on EOF.
|
||||||
|
//
|
||||||
|
// Verbosity::minimal — merges per-frame presence into contiguous time windows.
|
||||||
|
// Output: { "movie": "...", "actors": [{ "name", "imdb_id", "scenes": [[t0,t1], ...] }] }
|
||||||
|
//
|
||||||
|
// Verbosity::standard — per-frame detail including bboxes, similarity, unknowns.
|
||||||
|
// Output: { "frames": [{ "t", "identified": [...], "unknowns": [...] }] }
|
||||||
|
//
|
||||||
|
// eof signal: sets done_ = true so the main thread can call net.stop().
|
||||||
|
|
||||||
|
struct ResultSinkFunc {
|
||||||
|
static constexpr std::string_view label() { return "result_sink"; }
|
||||||
|
|
||||||
|
ResultSinkFunc(const Config& cfg, std::atomic<bool>& done)
|
||||||
|
: cfg_(cfg), done_(done)
|
||||||
|
{}
|
||||||
|
|
||||||
|
void operator()(SceneAnnotation sa) {
|
||||||
|
if (sa.eof) {
|
||||||
|
flush();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Progress to stderr
|
||||||
|
std::cerr << "\r[result_sink] t=" << sa.timestamp_sec << "s"
|
||||||
|
<< " active=" << count_known(sa.visible_actors)
|
||||||
|
<< " unknowns=" << count_unknown(sa.visible_actors)
|
||||||
|
<< std::flush;
|
||||||
|
|
||||||
|
frames_.push_back(std::move(sa));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Write accumulated results and signal done. Safe to call more than once.
|
||||||
|
void flush() {
|
||||||
|
if (written_.exchange(true)) return;
|
||||||
|
write_output();
|
||||||
|
done_.store(true, std::memory_order_release);
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
static int count_known(const std::vector<IdentifiedActor>& v) {
|
||||||
|
int n = 0;
|
||||||
|
for (const auto& a : v) if (a.actor_idx >= 0) ++n;
|
||||||
|
return n;
|
||||||
|
}
|
||||||
|
static int count_unknown(const std::vector<IdentifiedActor>& v) {
|
||||||
|
int n = 0;
|
||||||
|
for (const auto& a : v) if (a.actor_idx < 0) ++n;
|
||||||
|
return n;
|
||||||
|
}
|
||||||
|
|
||||||
|
void write_output() {
|
||||||
|
std::cerr << "\n[result_sink] writing " << cfg_.output_path << "\n";
|
||||||
|
|
||||||
|
json root;
|
||||||
|
if (cfg_.verbosity == Verbosity::xray) {
|
||||||
|
root = build_xray();
|
||||||
|
} else {
|
||||||
|
root["movie"] = cfg_.movie_path;
|
||||||
|
root["sample_fps"] = cfg_.sample_fps;
|
||||||
|
root["anneal_sec"] = cfg_.anneal_sec;
|
||||||
|
root["actors"] = build_epochs();
|
||||||
|
if (cfg_.verbosity == Verbosity::standard)
|
||||||
|
root["frames"] = build_standard();
|
||||||
|
}
|
||||||
|
|
||||||
|
std::ofstream f(cfg_.output_path);
|
||||||
|
if (!f.is_open()) {
|
||||||
|
std::cerr << "[result_sink] ERROR: cannot write " << cfg_.output_path << "\n";
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
f << root.dump(2) << "\n";
|
||||||
|
std::cerr << "[result_sink] done.\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ActorWindow {
|
||||||
|
std::string name, imdb_id;
|
||||||
|
std::vector<std::pair<double, double>> scenes; // [start_sec, end_sec]
|
||||||
|
};
|
||||||
|
|
||||||
|
// Core logic: merge per-frame detections into annealed [start, end] windows.
|
||||||
|
std::vector<ActorWindow> build_actor_windows() {
|
||||||
|
struct Info { std::string name, imdb_id; };
|
||||||
|
std::map<int, Info> actor_info;
|
||||||
|
std::map<int, std::vector<double>> timestamps;
|
||||||
|
|
||||||
|
for (const auto& frame : frames_) {
|
||||||
|
for (const auto& ia : frame.visible_actors) {
|
||||||
|
if (ia.actor_idx < 0) continue;
|
||||||
|
actor_info[ia.actor_idx] = {ia.name, ia.imdb_id};
|
||||||
|
timestamps[ia.actor_idx].push_back(frame.timestamp_sec);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<ActorWindow> result;
|
||||||
|
for (auto& [idx, ts_vec] : timestamps) {
|
||||||
|
ActorWindow aw;
|
||||||
|
aw.name = actor_info[idx].name;
|
||||||
|
aw.imdb_id = actor_info[idx].imdb_id;
|
||||||
|
|
||||||
|
double win_start = ts_vec[0], win_end = ts_vec[0];
|
||||||
|
for (size_t i = 1; i < ts_vec.size(); ++i) {
|
||||||
|
if (ts_vec[i] - win_end > cfg_.anneal_sec) {
|
||||||
|
aw.scenes.push_back({win_start, win_end});
|
||||||
|
win_start = ts_vec[i];
|
||||||
|
}
|
||||||
|
win_end = ts_vec[i];
|
||||||
|
}
|
||||||
|
aw.scenes.push_back({win_start, win_end});
|
||||||
|
result.push_back(std::move(aw));
|
||||||
|
}
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
json build_epochs() {
|
||||||
|
json actors = json::array();
|
||||||
|
for (const auto& aw : build_actor_windows()) {
|
||||||
|
json windows = json::array();
|
||||||
|
for (const auto& [s, e] : aw.scenes)
|
||||||
|
windows.push_back({s, e});
|
||||||
|
json ja;
|
||||||
|
ja["name"] = aw.name;
|
||||||
|
ja["imdb_id"] = aw.imdb_id;
|
||||||
|
ja["scenes"] = std::move(windows);
|
||||||
|
actors.push_back(std::move(ja));
|
||||||
|
}
|
||||||
|
return actors;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Jellyfin-Xray format: { "second": ["Actor", ...] }
|
||||||
|
// Expands each annealed window into every integer second so coverage is dense
|
||||||
|
// regardless of sample rate. Seconds between scenes have no key → overlay clears.
|
||||||
|
json build_xray() {
|
||||||
|
std::map<int, std::vector<std::string>> xray;
|
||||||
|
for (const auto& aw : build_actor_windows()) {
|
||||||
|
for (const auto& [start, end] : aw.scenes) {
|
||||||
|
int t0 = static_cast<int>(std::floor(start));
|
||||||
|
int t1 = static_cast<int>(std::ceil(end));
|
||||||
|
for (int t = t0; t <= t1; ++t)
|
||||||
|
xray[t].push_back(aw.name);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
json root = json::object();
|
||||||
|
for (const auto& [t, names] : xray)
|
||||||
|
root[std::to_string(t)] = names;
|
||||||
|
return root;
|
||||||
|
}
|
||||||
|
|
||||||
|
json build_standard() {
|
||||||
|
json frames = json::array();
|
||||||
|
for (const auto& frame : frames_) {
|
||||||
|
json jf;
|
||||||
|
jf["t"] = frame.timestamp_sec;
|
||||||
|
jf["identified"] = json::array();
|
||||||
|
jf["unknowns"] = json::array();
|
||||||
|
|
||||||
|
for (const auto& ia : frame.visible_actors) {
|
||||||
|
const auto& b = ia.bbox;
|
||||||
|
json jbox = {b.x, b.y, b.width, b.height};
|
||||||
|
|
||||||
|
if (ia.actor_idx >= 0) {
|
||||||
|
json ja;
|
||||||
|
ja["name"] = ia.name;
|
||||||
|
ja["imdb_id"] = ia.imdb_id;
|
||||||
|
ja["similarity"] = ia.similarity;
|
||||||
|
ja["track_id"] = ia.track_id;
|
||||||
|
ja["bbox"] = jbox;
|
||||||
|
jf["identified"].push_back(std::move(ja));
|
||||||
|
} else {
|
||||||
|
json ju;
|
||||||
|
ju["bbox"] = jbox;
|
||||||
|
ju["track_id"] = ia.track_id;
|
||||||
|
ju["confidence"] = ia.similarity; // reuse field; 0 for unknowns
|
||||||
|
jf["unknowns"].push_back(std::move(ju));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
frames.push_back(std::move(jf));
|
||||||
|
}
|
||||||
|
return frames;
|
||||||
|
}
|
||||||
|
|
||||||
|
const Config& cfg_;
|
||||||
|
std::atomic<bool>& done_;
|
||||||
|
std::atomic<bool> written_{false};
|
||||||
|
std::vector<SceneAnnotation> frames_;
|
||||||
|
};
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "config.hpp"
|
||||||
|
|
||||||
|
#include <map>
|
||||||
|
#include <iostream>
|
||||||
|
|
||||||
|
// ── SceneTrackerFunc ──────────────────────────────────────────────────────────
|
||||||
|
// KPN node: maintains an extinction-timer state machine per identified actor.
|
||||||
|
//
|
||||||
|
// On each MatchedSceneFrame:
|
||||||
|
// 1. Update last_seen for every matched known actor.
|
||||||
|
// 2. Expire actors whose last_seen is older than extinction_sec.
|
||||||
|
// 3. Emit SceneAnnotation with all currently active (non-expired) actors,
|
||||||
|
// including their most recently seen bbox and best similarity score.
|
||||||
|
//
|
||||||
|
// Unknown faces (actor_idx == -1) are passed through per-frame but are NOT
|
||||||
|
// tracked across frames — each frame reports its own unknowns independently.
|
||||||
|
|
||||||
|
struct SceneTrackerFunc {
|
||||||
|
static constexpr std::string_view label() { return "scene_tracker"; }
|
||||||
|
|
||||||
|
explicit SceneTrackerFunc(const Config& cfg)
|
||||||
|
: extinction_sec_(cfg.extinction_sec)
|
||||||
|
{
|
||||||
|
std::cerr << "[scene_tracker] extinction_sec=" << extinction_sec_ << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
SceneAnnotation operator()(MatchedSceneFrame mf) {
|
||||||
|
if (mf.source.eof) return {0.0, {}, /*eof=*/true};
|
||||||
|
|
||||||
|
double now = mf.source.timestamp_sec;
|
||||||
|
|
||||||
|
// Update known actors
|
||||||
|
for (const auto& ia : mf.actors) {
|
||||||
|
if (ia.actor_idx < 0) continue; // skip unknowns
|
||||||
|
|
||||||
|
auto& slot = active_[ia.actor_idx];
|
||||||
|
slot.last_seen = now;
|
||||||
|
slot.last_bbox = ia.bbox;
|
||||||
|
slot.last_crop = ia.crop;
|
||||||
|
slot.name = ia.name;
|
||||||
|
slot.imdb_id = ia.imdb_id;
|
||||||
|
// Keep the best (highest) similarity seen in this window
|
||||||
|
if (ia.similarity > slot.best_similarity)
|
||||||
|
slot.best_similarity = ia.similarity;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Expire stale actors
|
||||||
|
for (auto it = active_.begin(); it != active_.end(); ) {
|
||||||
|
if ((now - it->second.last_seen) > extinction_sec_)
|
||||||
|
it = active_.erase(it);
|
||||||
|
else
|
||||||
|
++it;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build annotation: active known actors
|
||||||
|
std::vector<IdentifiedActor> visible;
|
||||||
|
visible.reserve(active_.size() + mf.actors.size());
|
||||||
|
|
||||||
|
for (const auto& [actor_idx, slot] : active_) {
|
||||||
|
IdentifiedActor ia;
|
||||||
|
ia.actor_idx = actor_idx;
|
||||||
|
ia.name = slot.name;
|
||||||
|
ia.imdb_id = slot.imdb_id;
|
||||||
|
ia.similarity = slot.best_similarity;
|
||||||
|
ia.bbox = slot.last_bbox;
|
||||||
|
ia.crop = slot.last_crop;
|
||||||
|
visible.push_back(ia);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Append per-frame unknowns (actor_idx == -1) directly
|
||||||
|
for (const auto& ia : mf.actors) {
|
||||||
|
if (ia.actor_idx < 0) visible.push_back(ia);
|
||||||
|
}
|
||||||
|
|
||||||
|
return {now, std::move(visible)};
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
struct Slot {
|
||||||
|
double last_seen{0.0};
|
||||||
|
float best_similarity{0.f};
|
||||||
|
cv::Rect2f last_bbox;
|
||||||
|
cv::Mat last_crop;
|
||||||
|
std::string name;
|
||||||
|
std::string imdb_id;
|
||||||
|
};
|
||||||
|
|
||||||
|
double extinction_sec_;
|
||||||
|
std::map<int, Slot> active_; // actor_idx → state
|
||||||
|
};
|
||||||
@@ -0,0 +1,146 @@
|
|||||||
|
#pragma once
|
||||||
|
#include <onnxruntime/onnxruntime_cxx_api.h>
|
||||||
|
#include <filesystem>
|
||||||
|
#include <iostream>
|
||||||
|
#include <string>
|
||||||
|
#include <unordered_map>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// Detect the best available ORT execution provider and apply it to a
|
||||||
|
// SessionOptions. Priority order: TensorRT > CUDA > ROCm > CPU.
|
||||||
|
//
|
||||||
|
// Detection is conservative: GetAvailableProviders() confirms ORT was compiled
|
||||||
|
// with the provider, then AppendExecutionProvider_* is attempted inside a
|
||||||
|
// try/catch so a missing runtime library degrades gracefully to the next tier.
|
||||||
|
|
||||||
|
enum class OrtProvider { CPU, CUDA, ROCm, TensorRT };
|
||||||
|
|
||||||
|
// ── TRT configuration ─────────────────────────────────────────────────────────
|
||||||
|
// fp16: FP16 Tensor Core kernels — safe for both SCRFD and ArcFace.
|
||||||
|
// int8: INT8 quantisation — fast but UNSAFE for ArcFace without a
|
||||||
|
// calibration table (embedding cosine space will shift, breaking
|
||||||
|
// your similarity thresholds). Safe for the SCRFD detector.
|
||||||
|
// cache_dir: TRT engines are compiled once and cached here. First run is
|
||||||
|
// slow (~30–60 s per model); every subsequent run loads instantly.
|
||||||
|
// Shape profile (optional, set input_name + profile_{min,opt,max} to enable):
|
||||||
|
// ArcFace input is dynamic-batch (Nx3x112x112) — without a profile TRT
|
||||||
|
// builds at batch=1 and any larger call falls back to CUDA EP.
|
||||||
|
// SCRFD input is static-batch with dynamic H/W; we letterbox to 640×640
|
||||||
|
// and pin the profile to that.
|
||||||
|
// Shape strings are trtexec-style, e.g. "1x3x112x112".
|
||||||
|
|
||||||
|
struct TrtConfig {
|
||||||
|
bool fp16 = true;
|
||||||
|
bool int8 = false;
|
||||||
|
std::string cache_dir = "./trt_cache";
|
||||||
|
// Per-tensor optimisation profile. All four fields must be set together.
|
||||||
|
std::string input_name; // e.g. "input.1"
|
||||||
|
std::string profile_min; // e.g. "1x3x112x112"
|
||||||
|
std::string profile_opt; // e.g. "4x3x112x112"
|
||||||
|
std::string profile_max; // e.g. "8x3x112x112"
|
||||||
|
};
|
||||||
|
|
||||||
|
inline OrtProvider detect_ort_provider() {
|
||||||
|
auto available = Ort::GetAvailableProviders();
|
||||||
|
for (const auto& p : available) {
|
||||||
|
if (p == "TensorrtExecutionProvider") return OrtProvider::TensorRT;
|
||||||
|
if (p == "CUDAExecutionProvider") return OrtProvider::CUDA;
|
||||||
|
if (p == "ROCMExecutionProvider") return OrtProvider::ROCm;
|
||||||
|
}
|
||||||
|
return OrtProvider::CPU;
|
||||||
|
}
|
||||||
|
|
||||||
|
inline const char* provider_name(OrtProvider p) {
|
||||||
|
switch (p) {
|
||||||
|
case OrtProvider::TensorRT: return "TensorRT";
|
||||||
|
case OrtProvider::CUDA: return "CUDA";
|
||||||
|
case OrtProvider::ROCm: return "ROCm";
|
||||||
|
default: return "CPU";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply the given provider to opts. Falls back to CPU on failure and
|
||||||
|
// returns the provider that was actually applied.
|
||||||
|
inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
|
||||||
|
OrtProvider provider,
|
||||||
|
const char* label,
|
||||||
|
const TrtConfig& trt_cfg = {}) {
|
||||||
|
if (provider == OrtProvider::TensorRT) {
|
||||||
|
try {
|
||||||
|
std::filesystem::create_directories(trt_cfg.cache_dir);
|
||||||
|
|
||||||
|
// Use V2 API: key-value string map supports all options including
|
||||||
|
// dynamic batch profiles (missing from the legacy V1 struct).
|
||||||
|
std::unordered_map<std::string, std::string> kv = {
|
||||||
|
{"device_id", "0"},
|
||||||
|
{"trt_max_workspace_size", "2147483648"},
|
||||||
|
{"trt_fp16_enable", trt_cfg.fp16 ? "1" : "0"},
|
||||||
|
{"trt_int8_enable", trt_cfg.int8 ? "1" : "0"},
|
||||||
|
{"trt_engine_cache_enable", "1"},
|
||||||
|
{"trt_engine_cache_path", trt_cfg.cache_dir},
|
||||||
|
};
|
||||||
|
|
||||||
|
if (!trt_cfg.input_name.empty() && !trt_cfg.profile_min.empty()) {
|
||||||
|
kv["trt_profile_min_shapes"] =
|
||||||
|
trt_cfg.input_name + ":" + trt_cfg.profile_min;
|
||||||
|
kv["trt_profile_opt_shapes"] =
|
||||||
|
trt_cfg.input_name + ":" + trt_cfg.profile_opt;
|
||||||
|
kv["trt_profile_max_shapes"] =
|
||||||
|
trt_cfg.input_name + ":" + trt_cfg.profile_max;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ort::TensorRTProviderOptions trt_v2;
|
||||||
|
trt_v2.Update(kv);
|
||||||
|
opts.AppendExecutionProvider_TensorRT_V2(*trt_v2);
|
||||||
|
|
||||||
|
std::cerr << "[" << label << "] TensorRT"
|
||||||
|
<< (trt_cfg.fp16 ? " FP16" : "")
|
||||||
|
<< (trt_cfg.int8 ? " INT8" : "")
|
||||||
|
<< " cache=" << trt_cfg.cache_dir
|
||||||
|
<< (trt_cfg.profile_min.empty() ? "" :
|
||||||
|
" profile=" + trt_cfg.profile_min
|
||||||
|
+ "/" + trt_cfg.profile_opt
|
||||||
|
+ "/" + trt_cfg.profile_max)
|
||||||
|
<< "\n";
|
||||||
|
return OrtProvider::TensorRT;
|
||||||
|
} catch (const Ort::Exception& e) {
|
||||||
|
std::cerr << "[" << label << "] TensorRT unavailable ("
|
||||||
|
<< e.what() << "), trying CUDA\n";
|
||||||
|
provider = OrtProvider::CUDA;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (provider == OrtProvider::CUDA) {
|
||||||
|
try {
|
||||||
|
OrtCUDAProviderOptions cuda{};
|
||||||
|
cuda.device_id = 0;
|
||||||
|
opts.AppendExecutionProvider_CUDA(cuda);
|
||||||
|
std::cerr << "[" << label << "] CUDA provider\n";
|
||||||
|
return OrtProvider::CUDA;
|
||||||
|
} catch (const Ort::Exception& e) {
|
||||||
|
std::cerr << "[" << label << "] CUDA unavailable ("
|
||||||
|
<< e.what() << "), trying ROCm\n";
|
||||||
|
provider = OrtProvider::ROCm;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (provider == OrtProvider::ROCm) {
|
||||||
|
try {
|
||||||
|
OrtROCMProviderOptions rocm{};
|
||||||
|
rocm.device_id = 0;
|
||||||
|
opts.AppendExecutionProvider_ROCM(rocm);
|
||||||
|
std::cerr << "[" << label << "] ROCm provider\n";
|
||||||
|
return OrtProvider::ROCm;
|
||||||
|
} catch (const Ort::Exception& e) {
|
||||||
|
std::cerr << "[" << label << "] ROCm unavailable ("
|
||||||
|
<< e.what() << "), falling back to CPU\n";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
std::cerr << "[" << label << "] CPU provider\n";
|
||||||
|
return OrtProvider::CPU;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convenience: detect + apply in one call.
|
||||||
|
inline OrtProvider setup_ort_session(Ort::SessionOptions& opts,
|
||||||
|
const char* label,
|
||||||
|
const TrtConfig& trt_cfg = {}) {
|
||||||
|
return apply_ort_provider(opts, detect_ort_provider(), label, trt_cfg);
|
||||||
|
}
|
||||||
@@ -0,0 +1,165 @@
|
|||||||
|
// scene_preview — same pipeline as scene_analyze but with a live annotated
|
||||||
|
// display window driven from the main thread.
|
||||||
|
//
|
||||||
|
// KPN topology:
|
||||||
|
//
|
||||||
|
// [frame_source] ──► [face_detector] ──► [face_aligner] ──► [embedder]
|
||||||
|
// ──► [identity_matcher] ──► FanoutNode<MatchedSceneFrame,2>
|
||||||
|
// ├──► [scene_tracker] ──► [result_sink] (background thread)
|
||||||
|
// └──► [preview_node] (main thread)
|
||||||
|
//
|
||||||
|
// The main thread drives preview_node via preview.step(). When the movie ends
|
||||||
|
// (MatchedSceneFrame.source.eof == true) operator() returns false, the loop
|
||||||
|
// ends, and net.stop() is called. result_sink writes its JSON before its
|
||||||
|
// thread is joined by net.stop(), so the output file is always complete.
|
||||||
|
//
|
||||||
|
// Usage: same flags as scene_analyze (see main.cpp for reference).
|
||||||
|
// --preview-width <px> max display width in pixels (default: 1280)
|
||||||
|
|
||||||
|
#include "config.hpp"
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "types.hpp"
|
||||||
|
#include "gallery/gallery_store.hpp"
|
||||||
|
#include "nodes/frame_source_node.hpp"
|
||||||
|
#include "nodes/face_detector_node.hpp"
|
||||||
|
#include "nodes/face_aligner_node.hpp"
|
||||||
|
#include "nodes/embedder_node.hpp"
|
||||||
|
#include "nodes/face_tracker_node.hpp"
|
||||||
|
#include "nodes/identity_matcher_node.hpp"
|
||||||
|
#include "nodes/scene_tracker_node.hpp"
|
||||||
|
#include "nodes/result_sink_node.hpp"
|
||||||
|
#include "nodes/preview_node.hpp"
|
||||||
|
|
||||||
|
#include <kpn/kpn.hpp>
|
||||||
|
|
||||||
|
#include <atomic>
|
||||||
|
#include <chrono>
|
||||||
|
#include <cstring>
|
||||||
|
#include <iostream>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <thread>
|
||||||
|
|
||||||
|
static Config parse_args(int argc, char** argv) {
|
||||||
|
Config cfg;
|
||||||
|
cfg.detector_model = kDefaultDetectorModel;
|
||||||
|
cfg.arcface_model = kDefaultArcfaceModel;
|
||||||
|
cfg.output_path = "annotations.json";
|
||||||
|
cfg.verbosity = Verbosity::standard; // default to standard in preview mode
|
||||||
|
|
||||||
|
for (int i = 1; i < argc; ++i) {
|
||||||
|
auto arg = [&](const char* f) { return std::strcmp(argv[i], f) == 0; };
|
||||||
|
auto next = [&]() -> std::string {
|
||||||
|
if (++i >= argc) throw std::runtime_error(std::string("missing arg after ") + argv[i-1]);
|
||||||
|
return argv[i];
|
||||||
|
};
|
||||||
|
if (arg("--movie")) cfg.movie_path = next();
|
||||||
|
else if (arg("--gallery")) cfg.gallery_path = next();
|
||||||
|
else if (arg("--output")) cfg.output_path = next();
|
||||||
|
else if (arg("--fps")) cfg.sample_fps = std::stof(next());
|
||||||
|
else if (arg("--max-decode-fps")) cfg.max_decode_fps = std::stof(next());
|
||||||
|
else if (arg("--start")) cfg.start_sec = std::stod(next());
|
||||||
|
else if (arg("--end")) cfg.end_sec = std::stod(next());
|
||||||
|
else if (arg("--cut-threshold")) cfg.cut_threshold = std::stof(next());
|
||||||
|
else if (arg("--verbosity")) { int v = std::stoi(next()); cfg.verbosity = v == 2 ? Verbosity::xray : v == 1 ? Verbosity::standard : Verbosity::minimal; }
|
||||||
|
else if (arg("--prior")) cfg.match_prior = std::stof(next());
|
||||||
|
else if (arg("--prob-threshold")) cfg.prob_threshold = std::stof(next());
|
||||||
|
else if (arg("--match-threshold")) cfg.match_threshold = std::stof(next());
|
||||||
|
else if (arg("--extinction")) cfg.extinction_sec = std::stod(next());
|
||||||
|
else if (arg("--detector")) cfg.detector_model = next();
|
||||||
|
else if (arg("--detector-engine")) cfg.detector_engine = next();
|
||||||
|
else if (arg("--arcface")) cfg.arcface_model = next();
|
||||||
|
else if (arg("--arcface-engine")) cfg.arcface_engine = next();
|
||||||
|
else if (arg("--conf")) cfg.detector_conf = std::stof(next());
|
||||||
|
else if (arg("--max-faces")) cfg.max_faces = std::stoi(next());
|
||||||
|
else if (arg("--min-face-px")) cfg.min_face_px = std::stof(next());
|
||||||
|
else if (arg("--ratio")) cfg.match_ratio = std::stof(next());
|
||||||
|
else if (arg("--ratio-ceil")) cfg.match_ratio_ceil = std::stof(next());
|
||||||
|
else if (arg("--track-alpha")) cfg.track_alpha = std::stof(next());
|
||||||
|
else if (arg("--track-min-iou")) cfg.track_min_iou = std::stof(next());
|
||||||
|
else if (arg("--track-max-embed")) cfg.track_max_embed_dist = std::stof(next());
|
||||||
|
else if (arg("--track-max-missing")) cfg.track_max_frames_missing = std::stoi(next());
|
||||||
|
else if (arg("--track-min-frames")) cfg.track_min_frames = std::stoi(next());
|
||||||
|
else if (arg("--anneal")) cfg.anneal_sec = std::stod(next());
|
||||||
|
else if (arg("--trt-cache")) cfg.trt.cache_dir = next();
|
||||||
|
else if (arg("--trt-fp16")) cfg.trt.fp16 = true;
|
||||||
|
else if (arg("--no-trt-fp16")) cfg.trt.fp16 = false;
|
||||||
|
else if (arg("--trt-int8")) cfg.trt.int8 = true;
|
||||||
|
else if (arg("--embed-batch")) cfg.embed_batch_size = std::stoi(next());
|
||||||
|
else { std::cerr << "[warn] unknown flag: " << argv[i] << "\n"; }
|
||||||
|
}
|
||||||
|
if (cfg.movie_path.empty()) throw std::runtime_error("--movie is required");
|
||||||
|
if (cfg.gallery_path.empty()) throw std::runtime_error("--gallery is required");
|
||||||
|
return cfg;
|
||||||
|
}
|
||||||
|
|
||||||
|
int main(int argc, char** argv) {
|
||||||
|
Config cfg;
|
||||||
|
try { cfg = parse_args(argc, argv); }
|
||||||
|
catch (const std::exception& e) {
|
||||||
|
std::cerr << "Usage error: " << e.what() << "\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
ActorGallery gallery;
|
||||||
|
try { gallery = load_gallery(cfg.gallery_path); }
|
||||||
|
catch (const std::exception& e) {
|
||||||
|
std::cerr << "Gallery error: " << e.what() << "\n";
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Functors ──────────────────────────────────────────────────────────────
|
||||||
|
std::atomic<bool> done{false};
|
||||||
|
|
||||||
|
const OrtProvider provider = detect_ort_provider();
|
||||||
|
std::cerr << "[main] inference provider: " << provider_name(provider) << "\n";
|
||||||
|
|
||||||
|
FrameSourceFunc source_fn {cfg};
|
||||||
|
FaceDetectorFunc detector_fn{cfg, provider};
|
||||||
|
FaceAlignerFunc aligner_fn;
|
||||||
|
EmbedderFunc embedder_fn{cfg, provider};
|
||||||
|
FaceTrackerFunc ftracker_fn{cfg};
|
||||||
|
IdentityMatcherFunc matcher_fn {gallery, cfg};
|
||||||
|
SceneTrackerFunc tracker_fn {cfg};
|
||||||
|
ResultSinkFunc sink_fn {cfg, done};
|
||||||
|
|
||||||
|
// ── KPN ObjectNodes ───────────────────────────────────────────────────────
|
||||||
|
kpn::ObjectNode<FrameSourceFunc, kpn::in<>, kpn::out<"frame">, "frame_source", 0> source (source_fn, 32);
|
||||||
|
kpn::ObjectNode<FaceDetectorFunc, kpn::in<"frame">, kpn::out<"scene">, "face_detector", 0> detector (detector_fn, 64);
|
||||||
|
kpn::ObjectNode<FaceAlignerFunc, kpn::in<"scene">, kpn::out<"aligned">, "face_aligner", 0> aligner (aligner_fn, 64);
|
||||||
|
kpn::ObjectNode<EmbedderFunc, kpn::in<"aligned">, kpn::out<"embedded">, "embedder", 0> embedder (embedder_fn, 32);
|
||||||
|
kpn::ObjectNode<FaceTrackerFunc, kpn::in<"embedded">, kpn::out<"tracked">, "face_tracker", 0> ftracker (ftracker_fn, 16);
|
||||||
|
kpn::ObjectNode<IdentityMatcherFunc, kpn::in<"tracked">, kpn::out<"matched">, "identity_matcher", 0> matcher (matcher_fn, 16);
|
||||||
|
kpn::ObjectNode<SceneTrackerFunc, kpn::in<"matched">, kpn::out<"annotation">, "scene_tracker", 0> tracker (tracker_fn, 16);
|
||||||
|
kpn::ObjectNode<ResultSinkFunc, kpn::in<"annotation">,kpn::out<>, "result_sink", 0> sink (sink_fn, 16);
|
||||||
|
|
||||||
|
// MainThreadNode — no thread spawned; driven by preview.step() below
|
||||||
|
PreviewNode preview{cfg, 16};
|
||||||
|
|
||||||
|
// matcher → FanoutNode<MatchedSceneFrame,2> → [scene_tracker, preview] (auto-inserted)
|
||||||
|
auto net = kpn::make_network(
|
||||||
|
kpn::edge(source.output<"frame">(), detector.input<"frame">()),
|
||||||
|
kpn::edge(detector.output<"scene">(), aligner.input<"scene">()),
|
||||||
|
kpn::edge(aligner.output<"aligned">(), embedder.input<"aligned">()),
|
||||||
|
kpn::edge(embedder.output<"embedded">(), ftracker.input<"embedded">()),
|
||||||
|
kpn::edge(ftracker.output<"tracked">(), matcher.input<"tracked">()),
|
||||||
|
kpn::edge(matcher.output<"matched">(), tracker.input<"matched">()),
|
||||||
|
kpn::edge(matcher.output<"matched">(), preview.input<"matched">()),
|
||||||
|
kpn::edge(tracker.output<"annotation">(), sink.input<"annotation">())
|
||||||
|
);
|
||||||
|
|
||||||
|
// ── Run ───────────────────────────────────────────────────────────────────
|
||||||
|
std::cerr << "[main] starting pipeline — press q or Esc to quit early\n";
|
||||||
|
net.start();
|
||||||
|
|
||||||
|
// Main thread drives the display window; returns false on EOF or q/Esc
|
||||||
|
while (preview.step()) {
|
||||||
|
cv::waitKey(1); // pump OS events between frames
|
||||||
|
}
|
||||||
|
|
||||||
|
net.stop();
|
||||||
|
sink_fn.flush(); // write whatever was accumulated (no-op if EOF already flushed)
|
||||||
|
cv::destroyAllWindows();
|
||||||
|
net.print_diagnostics();
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
@@ -0,0 +1,230 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "ort_provider.hpp"
|
||||||
|
#include "types.hpp"
|
||||||
|
|
||||||
|
#include <onnxruntime/onnxruntime_cxx_api.h>
|
||||||
|
#include <opencv2/dnn.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// ── SCRFDDecoder ──────────────────────────────────────────────────────────────
|
||||||
|
// ONNX Runtime-based decoder for InsightFace SCRFD face detector with kps.
|
||||||
|
// Uses ORT instead of cv::dnn because OpenCV 4.x cannot load SCRFD's dynamic
|
||||||
|
// Shape nodes. ORT handles dynamic shapes natively and is thread-safe for
|
||||||
|
// concurrent Run() calls.
|
||||||
|
//
|
||||||
|
// Model output layout (9 tensors, InsightFace export order):
|
||||||
|
// [0-2] score_s8 / score_s16 / score_s32 — flat (N,)
|
||||||
|
// [3-5] bbox_s8 / bbox_s16 / bbox_s32 — flat (N*4,) distance format
|
||||||
|
// [6-8] kps_s8 / kps_s16 / kps_s32 — flat (N*10,) distance format
|
||||||
|
//
|
||||||
|
// Landmark order (same as YuNet/ArcFace convention):
|
||||||
|
// [0] right-eye [1] left-eye [2] nose [3] right-mouth [4] left-mouth
|
||||||
|
|
||||||
|
struct SCRFDDecoder {
|
||||||
|
static constexpr int kInputW = 640;
|
||||||
|
static constexpr int kInputH = 640;
|
||||||
|
static constexpr int kAllStrides[4] = {8, 16, 32, 64};
|
||||||
|
static constexpr int kAnchors = 2;
|
||||||
|
|
||||||
|
SCRFDDecoder(const std::string& model_path,
|
||||||
|
float conf_threshold, float nms_threshold,
|
||||||
|
OrtProvider provider = OrtProvider::CPU,
|
||||||
|
TrtConfig trt_cfg = {})
|
||||||
|
: conf_threshold_(conf_threshold)
|
||||||
|
, nms_threshold_(nms_threshold)
|
||||||
|
{
|
||||||
|
Ort::SessionOptions opts;
|
||||||
|
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||||
|
opts.SetIntraOpNumThreads(1);
|
||||||
|
|
||||||
|
// SCRFD ONNX has a dynamic H/W input; we letterbox to 640×640 at
|
||||||
|
// runtime, so pin the TRT profile to that single shape — otherwise
|
||||||
|
// TRT picks generic shapes and either rebuilds per-call or falls
|
||||||
|
// back to CUDA EP.
|
||||||
|
if (provider == OrtProvider::TensorRT) {
|
||||||
|
if (trt_cfg.input_name.empty()) {
|
||||||
|
Ort::SessionOptions probe_opts;
|
||||||
|
probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
|
||||||
|
Ort::Session probe(env_, model_path.c_str(), probe_opts);
|
||||||
|
Ort::AllocatorWithDefaultOptions alloc;
|
||||||
|
trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
|
||||||
|
}
|
||||||
|
if (trt_cfg.profile_min.empty()) {
|
||||||
|
const std::string shape =
|
||||||
|
"1x3x" + std::to_string(kInputH) + "x" + std::to_string(kInputW);
|
||||||
|
trt_cfg.profile_min = shape;
|
||||||
|
trt_cfg.profile_opt = shape;
|
||||||
|
trt_cfg.profile_max = shape;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
apply_ort_provider(opts, provider, "SCRFDDecoder", trt_cfg);
|
||||||
|
|
||||||
|
session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);
|
||||||
|
|
||||||
|
Ort::AllocatorWithDefaultOptions alloc;
|
||||||
|
auto in_name = session_->GetInputNameAllocated(0, alloc);
|
||||||
|
input_name_ = in_name.get();
|
||||||
|
|
||||||
|
const size_t n_out = session_->GetOutputCount();
|
||||||
|
if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
|
||||||
|
throw std::runtime_error(
|
||||||
|
"[SCRFDDecoder] expected 9 or 12 outputs (kps-variant model), got "
|
||||||
|
+ std::to_string(n_out));
|
||||||
|
fmc_ = static_cast<int>(n_out / 3);
|
||||||
|
|
||||||
|
for (size_t i = 0; i < n_out; ++i) {
|
||||||
|
auto name = session_->GetOutputNameAllocated(i, alloc);
|
||||||
|
out_name_storage_.emplace_back(name.get());
|
||||||
|
}
|
||||||
|
for (auto& s : out_name_storage_)
|
||||||
|
out_name_ptrs_.push_back(s.c_str());
|
||||||
|
|
||||||
|
// Reject non-SCRFD models (e.g. YuNet, which also has 12 outputs and so
|
||||||
|
// passes the count check above, but is encoded entirely differently).
|
||||||
|
// Cross-check by output channel count: SCRFD's three groups of fmc_
|
||||||
|
// outputs encode scores (1ch), bboxes (4ch) and 5-point kps (10ch).
|
||||||
|
// YuNet exports loc/conf/iou with 14/2/1 channels, so this trips.
|
||||||
|
const int expected_last[3] = {1, 4, 10};
|
||||||
|
for (size_t gi = 0; gi < 3; ++gi) {
|
||||||
|
for (int si = 0; si < fmc_; ++si) {
|
||||||
|
const size_t oi = gi * fmc_ + si;
|
||||||
|
auto shape = session_->GetOutputTypeInfo(oi)
|
||||||
|
.GetTensorTypeAndShapeInfo().GetShape();
|
||||||
|
if (shape.empty() || shape.back() != expected_last[gi]) {
|
||||||
|
throw std::runtime_error(
|
||||||
|
"[SCRFDDecoder] model does not look like InsightFace SCRFD: "
|
||||||
|
"output '" + out_name_storage_[oi] + "' last-dim is "
|
||||||
|
+ std::to_string(shape.empty() ? -1 : shape.back())
|
||||||
|
+ ", expected " + std::to_string(expected_last[gi])
|
||||||
|
+ ". Hint: pass scrfd_500m_bnkps.onnx, not yunet/*.onnx.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::cerr << "[SCRFDDecoder] loaded: " << model_path << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
// Thread-safe: ORT Run() is safe for concurrent calls on the same Session.
|
||||||
|
std::vector<DetectedFace> detect(const cv::Mat& img) const {
|
||||||
|
// Letterbox to 640×640: uniform scale (preserves aspect ratio) + pad
|
||||||
|
// shorter side with constant grey. Stretching to 640×640 (the prior
|
||||||
|
// behaviour) distorts faces non-uniformly and degrades landmark
|
||||||
|
// localisation — matters most for portrait gallery images and 16:9
|
||||||
|
// video frames alike. Coordinates are mapped back via inverse scale +
|
||||||
|
// pad-offset below.
|
||||||
|
const float scale = std::min(static_cast<float>(kInputW) / img.cols,
|
||||||
|
static_cast<float>(kInputH) / img.rows);
|
||||||
|
const int new_w = static_cast<int>(std::round(img.cols * scale));
|
||||||
|
const int new_h = static_cast<int>(std::round(img.rows * scale));
|
||||||
|
const int pad_x = (kInputW - new_w) / 2;
|
||||||
|
const int pad_y = (kInputH - new_h) / 2;
|
||||||
|
|
||||||
|
cv::Mat resized;
|
||||||
|
cv::resize(img, resized, {new_w, new_h}, 0, 0, cv::INTER_LINEAR);
|
||||||
|
cv::Mat letterboxed(kInputH, kInputW, img.type(),
|
||||||
|
cv::Scalar(114, 114, 114));
|
||||||
|
resized.copyTo(letterboxed(cv::Rect(pad_x, pad_y, new_w, new_h)));
|
||||||
|
|
||||||
|
// BGR→RGB swap + normalize to [-1,1] → NCHW float32 blob
|
||||||
|
cv::Mat blob = cv::dnn::blobFromImage(
|
||||||
|
letterboxed, 1.0 / 128.0, {kInputW, kInputH},
|
||||||
|
cv::Scalar(127.5f, 127.5f, 127.5f),
|
||||||
|
/*swapRB=*/true, /*crop=*/false, CV_32F);
|
||||||
|
|
||||||
|
const std::array<int64_t, 4> in_shape = {1, 3, kInputH, kInputW};
|
||||||
|
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||||||
|
auto in_tensor = Ort::Value::CreateTensor<float>(
|
||||||
|
mem, blob.ptr<float>(), blob.total(),
|
||||||
|
in_shape.data(), in_shape.size());
|
||||||
|
|
||||||
|
const char* in_name_c = input_name_.c_str();
|
||||||
|
auto outs = session_->Run(
|
||||||
|
Ort::RunOptions{nullptr},
|
||||||
|
&in_name_c, &in_tensor, 1,
|
||||||
|
out_name_ptrs_.data(), out_name_ptrs_.size());
|
||||||
|
|
||||||
|
std::vector<cv::Rect2d> raw_boxes;
|
||||||
|
std::vector<float> raw_scores;
|
||||||
|
std::vector<std::array<cv::Point2f, 5>> raw_kps;
|
||||||
|
|
||||||
|
for (int si = 0; si < fmc_; ++si) {
|
||||||
|
const int stride = kAllStrides[si];
|
||||||
|
const int fh = kInputH / stride;
|
||||||
|
const int fw = kInputW / stride;
|
||||||
|
|
||||||
|
const float* s = outs[si].GetTensorData<float>();
|
||||||
|
const float* b = outs[fmc_ + si].GetTensorData<float>();
|
||||||
|
const float* k = outs[fmc_ * 2 + si].GetTensorData<float>();
|
||||||
|
|
||||||
|
for (int r = 0; r < fh; ++r) {
|
||||||
|
for (int c = 0; c < fw; ++c) {
|
||||||
|
for (int a = 0; a < kAnchors; ++a) {
|
||||||
|
const int idx = (r * fw + c) * kAnchors + a;
|
||||||
|
const float score = s[idx];
|
||||||
|
if (score < conf_threshold_) continue;
|
||||||
|
|
||||||
|
const float cx = static_cast<float>(c * stride);
|
||||||
|
const float cy = static_cast<float>(r * stride);
|
||||||
|
|
||||||
|
// Decode in letterboxed network space, then un-pad +
|
||||||
|
// un-scale to original image coordinates.
|
||||||
|
const auto to_img_x = [&](float v) { return (v - pad_x) / scale; };
|
||||||
|
const auto to_img_y = [&](float v) { return (v - pad_y) / scale; };
|
||||||
|
|
||||||
|
const float x1 = to_img_x(cx - b[idx*4+0] * stride);
|
||||||
|
const float y1 = to_img_y(cy - b[idx*4+1] * stride);
|
||||||
|
const float x2 = to_img_x(cx + b[idx*4+2] * stride);
|
||||||
|
const float y2 = to_img_y(cy + b[idx*4+3] * stride);
|
||||||
|
raw_boxes.push_back({(double)x1, (double)y1,
|
||||||
|
(double)(x2-x1), (double)(y2-y1)});
|
||||||
|
raw_scores.push_back(score);
|
||||||
|
|
||||||
|
std::array<cv::Point2f, 5> lms;
|
||||||
|
for (int p = 0; p < 5; ++p)
|
||||||
|
lms[p] = {to_img_x(cx + k[idx*10+p*2 ] * stride),
|
||||||
|
to_img_y(cy + k[idx*10+p*2+1] * stride)};
|
||||||
|
raw_kps.push_back(lms);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<int> keep;
|
||||||
|
cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep);
|
||||||
|
|
||||||
|
const float img_w = static_cast<float>(img.cols);
|
||||||
|
const float img_h = static_cast<float>(img.rows);
|
||||||
|
|
||||||
|
std::vector<DetectedFace> faces;
|
||||||
|
faces.reserve(keep.size());
|
||||||
|
for (int i : keep) {
|
||||||
|
const auto& rb = raw_boxes[i];
|
||||||
|
DetectedFace f;
|
||||||
|
const float x = std::max(0.f, (float)rb.x);
|
||||||
|
const float y = std::max(0.f, (float)rb.y);
|
||||||
|
f.bbox = {x, y,
|
||||||
|
std::min((float)rb.width, img_w - x),
|
||||||
|
std::min((float)rb.height, img_h - y)};
|
||||||
|
f.confidence = raw_scores[i];
|
||||||
|
f.landmarks = raw_kps[i];
|
||||||
|
faces.push_back(f);
|
||||||
|
}
|
||||||
|
return faces;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
float conf_threshold_;
|
||||||
|
float nms_threshold_;
|
||||||
|
int fmc_{3};
|
||||||
|
Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "scrfd"};
|
||||||
|
std::unique_ptr<Ort::Session> session_;
|
||||||
|
std::string input_name_;
|
||||||
|
std::vector<std::string> out_name_storage_;
|
||||||
|
std::vector<const char*> out_name_ptrs_;
|
||||||
|
};
|
||||||
@@ -0,0 +1,202 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "face_utils.hpp"
|
||||||
|
#include "types.hpp"
|
||||||
|
|
||||||
|
#include <NvInfer.h>
|
||||||
|
#include <cuda_runtime_api.h>
|
||||||
|
#include <opencv2/dnn.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <cstdint>
|
||||||
|
#include <cstring>
|
||||||
|
#include <fstream>
|
||||||
|
#include <iostream>
|
||||||
|
#include <memory>
|
||||||
|
#include <mutex>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// ── TrtArcFaceEmbedder ────────────────────────────────────────────────────────
|
||||||
|
// Pure-TensorRT ArcFace runner. Loads a serialised engine built by
|
||||||
|
// scripts/build_trt_engines.sh (or any trtexec-produced .engine matching the
|
||||||
|
// ArcFace I/O contract: input Nx3x112x112 float32, output Nx512 float32 or
|
||||||
|
// float16).
|
||||||
|
//
|
||||||
|
// Skips ONNX Runtime entirely — useful on systems where ORT was built without
|
||||||
|
// the TensorRT EP (e.g. Arch's onnxruntime-opt-cuda 1.24.x).
|
||||||
|
//
|
||||||
|
// Thread-safety: a single IExecutionContext is not safe to drive from multiple
|
||||||
|
// threads concurrently; we serialise with a mutex. The KPN embedder node is
|
||||||
|
// single-threaded anyway.
|
||||||
|
|
||||||
|
namespace trt_arcface_detail {
|
||||||
|
inline std::string trim_path(const std::string& s) { return s; }
|
||||||
|
|
||||||
|
struct CudaError : std::runtime_error {
|
||||||
|
using std::runtime_error::runtime_error;
|
||||||
|
};
|
||||||
|
|
||||||
|
inline void check_cuda(cudaError_t e, const char* what) {
|
||||||
|
if (e != cudaSuccess)
|
||||||
|
throw CudaError(std::string(what) + ": " + cudaGetErrorString(e));
|
||||||
|
}
|
||||||
|
|
||||||
|
class TrtLogger : public nvinfer1::ILogger {
|
||||||
|
public:
|
||||||
|
void log(Severity sev, const char* msg) noexcept override {
|
||||||
|
if (sev <= Severity::kWARNING)
|
||||||
|
std::cerr << "[TRT] " << msg << "\n";
|
||||||
|
}
|
||||||
|
};
|
||||||
|
inline TrtLogger& logger() { static TrtLogger g; return g; }
|
||||||
|
} // namespace trt_arcface_detail
|
||||||
|
|
||||||
|
struct TrtArcFaceEmbedder {
|
||||||
|
explicit TrtArcFaceEmbedder(const std::string& engine_path) {
|
||||||
|
using namespace trt_arcface_detail;
|
||||||
|
std::ifstream f(engine_path, std::ios::binary | std::ios::ate);
|
||||||
|
if (!f) throw std::runtime_error("TrtArcFaceEmbedder: cannot open " + engine_path);
|
||||||
|
const std::streamsize sz = f.tellg();
|
||||||
|
f.seekg(0);
|
||||||
|
std::vector<char> blob(sz);
|
||||||
|
f.read(blob.data(), sz);
|
||||||
|
|
||||||
|
runtime_.reset(nvinfer1::createInferRuntime(logger()));
|
||||||
|
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
||||||
|
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), sz));
|
||||||
|
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
|
||||||
|
context_.reset(engine_->createExecutionContext());
|
||||||
|
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
||||||
|
|
||||||
|
// Resolve I/O tensor names + the max batch the profile permits.
|
||||||
|
const int n_io = engine_->getNbIOTensors();
|
||||||
|
for (int i = 0; i < n_io; ++i) {
|
||||||
|
const char* name = engine_->getIOTensorName(i);
|
||||||
|
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)
|
||||||
|
input_name_ = name;
|
||||||
|
else
|
||||||
|
output_name_ = name;
|
||||||
|
}
|
||||||
|
if (input_name_.empty() || output_name_.empty())
|
||||||
|
throw std::runtime_error("TrtArcFaceEmbedder: engine missing input/output tensor");
|
||||||
|
|
||||||
|
auto in_dtype = engine_->getTensorDataType(input_name_.c_str());
|
||||||
|
auto out_dtype = engine_->getTensorDataType(output_name_.c_str());
|
||||||
|
input_is_fp16_ = (in_dtype == nvinfer1::DataType::kHALF);
|
||||||
|
output_is_fp16_ = (out_dtype == nvinfer1::DataType::kHALF);
|
||||||
|
|
||||||
|
auto max_dims = engine_->getProfileShape(input_name_.c_str(), 0,
|
||||||
|
nvinfer1::OptProfileSelector::kMAX);
|
||||||
|
if (max_dims.nbDims != 4 || max_dims.d[1] != 3 ||
|
||||||
|
max_dims.d[2] != 112 || max_dims.d[3] != 112)
|
||||||
|
throw std::runtime_error("TrtArcFaceEmbedder: unexpected input shape in engine");
|
||||||
|
max_batch_ = max_dims.d[0];
|
||||||
|
|
||||||
|
const std::size_t in_bytes = static_cast<std::size_t>(max_batch_) * 3 * 112 * 112 *
|
||||||
|
(input_is_fp16_ ? 2 : 4);
|
||||||
|
const std::size_t out_bytes = static_cast<std::size_t>(max_batch_) * 512 *
|
||||||
|
(output_is_fp16_ ? 2 : 4);
|
||||||
|
check_cuda(cudaMalloc(&d_input_, in_bytes), "cudaMalloc input");
|
||||||
|
check_cuda(cudaMalloc(&d_output_, out_bytes), "cudaMalloc output");
|
||||||
|
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
|
||||||
|
|
||||||
|
context_->setTensorAddress(input_name_.c_str(), d_input_);
|
||||||
|
context_->setTensorAddress(output_name_.c_str(), d_output_);
|
||||||
|
|
||||||
|
std::cerr << "[TrtArcFace] loaded: " << engine_path
|
||||||
|
<< " max_batch=" << max_batch_
|
||||||
|
<< (input_is_fp16_ ? " fp16-in" : "")
|
||||||
|
<< (output_is_fp16_ ? " fp16-out" : "")
|
||||||
|
<< "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
~TrtArcFaceEmbedder() {
|
||||||
|
if (stream_) cudaStreamDestroy(stream_);
|
||||||
|
if (d_input_) cudaFree(d_input_);
|
||||||
|
if (d_output_) cudaFree(d_output_);
|
||||||
|
}
|
||||||
|
|
||||||
|
TrtArcFaceEmbedder(const TrtArcFaceEmbedder&) = delete;
|
||||||
|
TrtArcFaceEmbedder& operator=(const TrtArcFaceEmbedder&) = delete;
|
||||||
|
|
||||||
|
int max_batch() const { return max_batch_; }
|
||||||
|
|
||||||
|
std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) const {
|
||||||
|
using namespace trt_arcface_detail;
|
||||||
|
if (crops.empty()) return {};
|
||||||
|
const int n = static_cast<int>(crops.size());
|
||||||
|
if (n > max_batch_)
|
||||||
|
throw std::runtime_error("TrtArcFaceEmbedder: batch " + std::to_string(n) +
|
||||||
|
" exceeds engine max " + std::to_string(max_batch_));
|
||||||
|
|
||||||
|
std::vector<cv::Mat> rgbs(n);
|
||||||
|
for (int i = 0; i < n; ++i)
|
||||||
|
cv::cvtColor(crops[i], rgbs[i], cv::COLOR_BGR2RGB);
|
||||||
|
cv::Mat blob = cv::dnn::blobFromImages(
|
||||||
|
rgbs, 1.0 / 128.0, {112, 112},
|
||||||
|
cv::Scalar(127.5, 127.5, 127.5),
|
||||||
|
/*swapRB=*/false, /*crop=*/false, CV_32F);
|
||||||
|
|
||||||
|
std::lock_guard<std::mutex> lk(mu_);
|
||||||
|
context_->setInputShape(input_name_.c_str(),
|
||||||
|
nvinfer1::Dims4{n, 3, 112, 112});
|
||||||
|
|
||||||
|
const std::size_t in_count = static_cast<std::size_t>(n) * 3 * 112 * 112;
|
||||||
|
if (input_is_fp16_) {
|
||||||
|
cv::Mat blob16;
|
||||||
|
blob.convertTo(blob16, CV_16F);
|
||||||
|
check_cuda(cudaMemcpyAsync(d_input_, blob16.ptr(), in_count * 2,
|
||||||
|
cudaMemcpyHostToDevice, stream_),
|
||||||
|
"H2D input fp16");
|
||||||
|
} else {
|
||||||
|
check_cuda(cudaMemcpyAsync(d_input_, blob.ptr<float>(), in_count * 4,
|
||||||
|
cudaMemcpyHostToDevice, stream_),
|
||||||
|
"H2D input fp32");
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!context_->enqueueV3(stream_))
|
||||||
|
throw std::runtime_error("TrtArcFaceEmbedder: enqueueV3 failed");
|
||||||
|
|
||||||
|
const std::size_t out_count = static_cast<std::size_t>(n) * 512;
|
||||||
|
std::vector<float> host_f32(out_count);
|
||||||
|
if (output_is_fp16_) {
|
||||||
|
std::vector<uint16_t> host_f16(out_count);
|
||||||
|
check_cuda(cudaMemcpyAsync(host_f16.data(), d_output_, out_count * 2,
|
||||||
|
cudaMemcpyDeviceToHost, stream_),
|
||||||
|
"D2H output fp16");
|
||||||
|
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||||||
|
cv::Mat src16(1, static_cast<int>(out_count), CV_16F, host_f16.data());
|
||||||
|
cv::Mat dst32(1, static_cast<int>(out_count), CV_32F, host_f32.data());
|
||||||
|
src16.convertTo(dst32, CV_32F);
|
||||||
|
} else {
|
||||||
|
check_cuda(cudaMemcpyAsync(host_f32.data(), d_output_, out_count * 4,
|
||||||
|
cudaMemcpyDeviceToHost, stream_),
|
||||||
|
"D2H output fp32");
|
||||||
|
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<Embedding> out(n);
|
||||||
|
for (int i = 0; i < n; ++i)
|
||||||
|
out[i] = l2_normalise(host_f32.data() + i * 512);
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
struct TrtDeleter { template<class T> void operator()(T* p) const { delete p; } };
|
||||||
|
std::unique_ptr<nvinfer1::IRuntime, TrtDeleter> runtime_;
|
||||||
|
std::unique_ptr<nvinfer1::ICudaEngine, TrtDeleter> engine_;
|
||||||
|
std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
|
||||||
|
|
||||||
|
std::string input_name_;
|
||||||
|
std::string output_name_;
|
||||||
|
bool input_is_fp16_ = false;
|
||||||
|
bool output_is_fp16_ = false;
|
||||||
|
int max_batch_ = 1;
|
||||||
|
|
||||||
|
void* d_input_ = nullptr;
|
||||||
|
void* d_output_ = nullptr;
|
||||||
|
cudaStream_t stream_ = nullptr;
|
||||||
|
|
||||||
|
mutable std::mutex mu_;
|
||||||
|
};
|
||||||
@@ -0,0 +1,263 @@
|
|||||||
|
#pragma once
|
||||||
|
#include "trt_arcface_embedder.hpp" // pulls in CudaError/check_cuda/TrtLogger + nvinfer1/cuda headers
|
||||||
|
#include "types.hpp"
|
||||||
|
|
||||||
|
#include <opencv2/dnn.hpp>
|
||||||
|
#include <opencv2/imgproc.hpp>
|
||||||
|
|
||||||
|
#include <algorithm>
|
||||||
|
#include <array>
|
||||||
|
#include <cstring>
|
||||||
|
#include <fstream>
|
||||||
|
#include <iostream>
|
||||||
|
#include <memory>
|
||||||
|
#include <mutex>
|
||||||
|
#include <stdexcept>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
// ── TrtScrfdDecoder ───────────────────────────────────────────────────────────
|
||||||
|
// Pure-TensorRT SCRFD face detector. Loads a serialised engine built by
|
||||||
|
// scripts/build_trt_engines.sh (1x3x640x640 input pinned). Post-processing
|
||||||
|
// matches scrfd_decoder.hpp byte-for-byte — only inference is swapped.
|
||||||
|
//
|
||||||
|
// Output layout (9 tensors, InsightFace export order — preserved by trtexec):
|
||||||
|
// [0..2] score_s8 / score_s16 / score_s32 (N,1)
|
||||||
|
// [3..5] bbox_s8 / bbox_s16 / bbox_s32 (N,4)
|
||||||
|
// [6..8] kps_s8 / kps_s16 / kps_s32 (N,10)
|
||||||
|
|
||||||
|
struct TrtScrfdDecoder {
|
||||||
|
static constexpr int kInputW = 640;
|
||||||
|
static constexpr int kInputH = 640;
|
||||||
|
static constexpr int kAllStrides[4] = {8, 16, 32, 64};
|
||||||
|
static constexpr int kAnchors = 2;
|
||||||
|
|
||||||
|
TrtScrfdDecoder(const std::string& engine_path,
|
||||||
|
float conf_threshold, float nms_threshold)
|
||||||
|
: conf_threshold_(conf_threshold)
|
||||||
|
, nms_threshold_(nms_threshold)
|
||||||
|
{
|
||||||
|
using namespace trt_arcface_detail;
|
||||||
|
std::ifstream f(engine_path, std::ios::binary | std::ios::ate);
|
||||||
|
if (!f) throw std::runtime_error("TrtScrfdDecoder: cannot open " + engine_path);
|
||||||
|
const std::streamsize sz = f.tellg();
|
||||||
|
f.seekg(0);
|
||||||
|
std::vector<char> blob(sz);
|
||||||
|
f.read(blob.data(), sz);
|
||||||
|
|
||||||
|
runtime_.reset(nvinfer1::createInferRuntime(logger()));
|
||||||
|
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
||||||
|
engine_.reset(runtime_->deserializeCudaEngine(blob.data(), sz));
|
||||||
|
if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
|
||||||
|
context_.reset(engine_->createExecutionContext());
|
||||||
|
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
||||||
|
|
||||||
|
// Enumerate I/O tensors preserving engine declaration order.
|
||||||
|
const int n_io = engine_->getNbIOTensors();
|
||||||
|
for (int i = 0; i < n_io; ++i) {
|
||||||
|
const char* name = engine_->getIOTensorName(i);
|
||||||
|
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) {
|
||||||
|
if (!input_name_.empty())
|
||||||
|
throw std::runtime_error("TrtScrfdDecoder: multiple inputs not supported");
|
||||||
|
input_name_ = name;
|
||||||
|
} else {
|
||||||
|
output_names_.emplace_back(name);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (input_name_.empty())
|
||||||
|
throw std::runtime_error("TrtScrfdDecoder: no input tensor");
|
||||||
|
const int n_out = static_cast<int>(output_names_.size());
|
||||||
|
if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
|
||||||
|
throw std::runtime_error(
|
||||||
|
"TrtScrfdDecoder: expected 9 or 12 outputs (kps-variant SCRFD), got "
|
||||||
|
+ std::to_string(n_out));
|
||||||
|
fmc_ = n_out / 3;
|
||||||
|
|
||||||
|
// Validate input shape; the engine was built with min=opt=max=1x3x640x640.
|
||||||
|
auto in_dims = engine_->getProfileShape(input_name_.c_str(), 0,
|
||||||
|
nvinfer1::OptProfileSelector::kOPT);
|
||||||
|
if (in_dims.nbDims != 4 || in_dims.d[0] != 1 || in_dims.d[1] != 3 ||
|
||||||
|
in_dims.d[2] != kInputH || in_dims.d[3] != kInputW)
|
||||||
|
throw std::runtime_error(
|
||||||
|
"TrtScrfdDecoder: engine input must be 1x3x" +
|
||||||
|
std::to_string(kInputH) + "x" + std::to_string(kInputW));
|
||||||
|
|
||||||
|
// Allocate device buffer for input.
|
||||||
|
const std::size_t in_bytes = static_cast<std::size_t>(3) * kInputH * kInputW * 4;
|
||||||
|
check_cuda(cudaMalloc(&d_input_, in_bytes), "cudaMalloc input");
|
||||||
|
context_->setTensorAddress(input_name_.c_str(), d_input_);
|
||||||
|
context_->setInputShape(input_name_.c_str(),
|
||||||
|
nvinfer1::Dims4{1, 3, kInputH, kInputW});
|
||||||
|
|
||||||
|
// Allocate device + host buffers for each output, sized from engine.
|
||||||
|
d_outputs_.resize(n_out, nullptr);
|
||||||
|
host_outputs_.resize(n_out);
|
||||||
|
out_elem_counts_.resize(n_out, 0);
|
||||||
|
out_last_dims_.resize(n_out, 0);
|
||||||
|
|
||||||
|
const int expected_last[3] = {1, 4, 10};
|
||||||
|
for (int oi = 0; oi < n_out; ++oi) {
|
||||||
|
auto dims = context_->getTensorShape(output_names_[oi].c_str());
|
||||||
|
if (dims.nbDims < 1)
|
||||||
|
throw std::runtime_error("TrtScrfdDecoder: bad shape for output " +
|
||||||
|
output_names_[oi]);
|
||||||
|
std::size_t count = 1;
|
||||||
|
for (int d = 0; d < dims.nbDims; ++d) count *= static_cast<std::size_t>(dims.d[d]);
|
||||||
|
const int last = dims.d[dims.nbDims - 1];
|
||||||
|
const int group = oi / fmc_; // 0=scores, 1=bboxes, 2=kps
|
||||||
|
if (last != expected_last[group])
|
||||||
|
throw std::runtime_error(
|
||||||
|
"TrtScrfdDecoder: output '" + output_names_[oi] + "' last-dim is " +
|
||||||
|
std::to_string(last) + ", expected " + std::to_string(expected_last[group]) +
|
||||||
|
". Engine does not match SCRFD-bnkps layout.");
|
||||||
|
|
||||||
|
check_cuda(cudaMalloc(&d_outputs_[oi], count * 4), "cudaMalloc output");
|
||||||
|
context_->setTensorAddress(output_names_[oi].c_str(), d_outputs_[oi]);
|
||||||
|
host_outputs_[oi].resize(count);
|
||||||
|
out_elem_counts_[oi] = count;
|
||||||
|
out_last_dims_[oi] = last;
|
||||||
|
}
|
||||||
|
|
||||||
|
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
|
||||||
|
|
||||||
|
std::cerr << "[TrtScrfd] loaded: " << engine_path
|
||||||
|
<< " fmc=" << fmc_
|
||||||
|
<< " outputs=" << n_out << "\n";
|
||||||
|
}
|
||||||
|
|
||||||
|
~TrtScrfdDecoder() {
|
||||||
|
if (stream_) cudaStreamDestroy(stream_);
|
||||||
|
if (d_input_) cudaFree(d_input_);
|
||||||
|
for (void* p : d_outputs_) if (p) cudaFree(p);
|
||||||
|
}
|
||||||
|
|
||||||
|
TrtScrfdDecoder(const TrtScrfdDecoder&) = delete;
|
||||||
|
TrtScrfdDecoder& operator=(const TrtScrfdDecoder&) = delete;
|
||||||
|
|
||||||
|
std::vector<DetectedFace> detect(const cv::Mat& img) const {
|
||||||
|
using namespace trt_arcface_detail;
|
||||||
|
|
||||||
|
// Letterbox to 640×640 — identical to SCRFDDecoder.
|
||||||
|
const float scale = std::min(static_cast<float>(kInputW) / img.cols,
|
||||||
|
static_cast<float>(kInputH) / img.rows);
|
||||||
|
const int new_w = static_cast<int>(std::round(img.cols * scale));
|
||||||
|
const int new_h = static_cast<int>(std::round(img.rows * scale));
|
||||||
|
const int pad_x = (kInputW - new_w) / 2;
|
||||||
|
const int pad_y = (kInputH - new_h) / 2;
|
||||||
|
|
||||||
|
cv::Mat resized;
|
||||||
|
cv::resize(img, resized, {new_w, new_h}, 0, 0, cv::INTER_LINEAR);
|
||||||
|
cv::Mat letterboxed(kInputH, kInputW, img.type(), cv::Scalar(114, 114, 114));
|
||||||
|
resized.copyTo(letterboxed(cv::Rect(pad_x, pad_y, new_w, new_h)));
|
||||||
|
|
||||||
|
cv::Mat blob = cv::dnn::blobFromImage(
|
||||||
|
letterboxed, 1.0 / 128.0, {kInputW, kInputH},
|
||||||
|
cv::Scalar(127.5f, 127.5f, 127.5f),
|
||||||
|
/*swapRB=*/true, /*crop=*/false, CV_32F);
|
||||||
|
|
||||||
|
std::lock_guard<std::mutex> lk(mu_);
|
||||||
|
const std::size_t in_count = static_cast<std::size_t>(3) * kInputH * kInputW;
|
||||||
|
check_cuda(cudaMemcpyAsync(d_input_, blob.ptr<float>(), in_count * 4,
|
||||||
|
cudaMemcpyHostToDevice, stream_),
|
||||||
|
"H2D input");
|
||||||
|
|
||||||
|
if (!context_->enqueueV3(stream_))
|
||||||
|
throw std::runtime_error("TrtScrfdDecoder: enqueueV3 failed");
|
||||||
|
|
||||||
|
for (std::size_t oi = 0; oi < d_outputs_.size(); ++oi) {
|
||||||
|
check_cuda(cudaMemcpyAsync(host_outputs_[oi].data(), d_outputs_[oi],
|
||||||
|
out_elem_counts_[oi] * 4,
|
||||||
|
cudaMemcpyDeviceToHost, stream_),
|
||||||
|
"D2H output");
|
||||||
|
}
|
||||||
|
check_cuda(cudaStreamSynchronize(stream_), "stream sync");
|
||||||
|
|
||||||
|
// ── Post-process: identical to SCRFDDecoder ───────────────────────────
|
||||||
|
std::vector<cv::Rect2d> raw_boxes;
|
||||||
|
std::vector<float> raw_scores;
|
||||||
|
std::vector<std::array<cv::Point2f, 5>> raw_kps;
|
||||||
|
|
||||||
|
for (int si = 0; si < fmc_; ++si) {
|
||||||
|
const int stride = kAllStrides[si];
|
||||||
|
const int fh = kInputH / stride;
|
||||||
|
const int fw = kInputW / stride;
|
||||||
|
|
||||||
|
const float* s = host_outputs_[si].data();
|
||||||
|
const float* b = host_outputs_[fmc_ + si].data();
|
||||||
|
const float* k = host_outputs_[fmc_ * 2 + si].data();
|
||||||
|
|
||||||
|
for (int r = 0; r < fh; ++r) {
|
||||||
|
for (int c = 0; c < fw; ++c) {
|
||||||
|
for (int a = 0; a < kAnchors; ++a) {
|
||||||
|
const int idx = (r * fw + c) * kAnchors + a;
|
||||||
|
const float score = s[idx];
|
||||||
|
if (score < conf_threshold_) continue;
|
||||||
|
|
||||||
|
const float cx = static_cast<float>(c * stride);
|
||||||
|
const float cy = static_cast<float>(r * stride);
|
||||||
|
|
||||||
|
const auto to_img_x = [&](float v) { return (v - pad_x) / scale; };
|
||||||
|
const auto to_img_y = [&](float v) { return (v - pad_y) / scale; };
|
||||||
|
|
||||||
|
const float x1 = to_img_x(cx - b[idx*4+0] * stride);
|
||||||
|
const float y1 = to_img_y(cy - b[idx*4+1] * stride);
|
||||||
|
const float x2 = to_img_x(cx + b[idx*4+2] * stride);
|
||||||
|
const float y2 = to_img_y(cy + b[idx*4+3] * stride);
|
||||||
|
raw_boxes.push_back({(double)x1, (double)y1,
|
||||||
|
(double)(x2-x1), (double)(y2-y1)});
|
||||||
|
raw_scores.push_back(score);
|
||||||
|
|
||||||
|
std::array<cv::Point2f, 5> lms;
|
||||||
|
for (int p = 0; p < 5; ++p)
|
||||||
|
lms[p] = {to_img_x(cx + k[idx*10+p*2 ] * stride),
|
||||||
|
to_img_y(cy + k[idx*10+p*2+1] * stride)};
|
||||||
|
raw_kps.push_back(lms);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<int> keep;
|
||||||
|
cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep);
|
||||||
|
|
||||||
|
const float img_w = static_cast<float>(img.cols);
|
||||||
|
const float img_h = static_cast<float>(img.rows);
|
||||||
|
|
||||||
|
std::vector<DetectedFace> faces;
|
||||||
|
faces.reserve(keep.size());
|
||||||
|
for (int i : keep) {
|
||||||
|
const auto& rb = raw_boxes[i];
|
||||||
|
DetectedFace f;
|
||||||
|
const float x = std::max(0.f, (float)rb.x);
|
||||||
|
const float y = std::max(0.f, (float)rb.y);
|
||||||
|
f.bbox = {x, y,
|
||||||
|
std::min((float)rb.width, img_w - x),
|
||||||
|
std::min((float)rb.height, img_h - y)};
|
||||||
|
f.confidence = raw_scores[i];
|
||||||
|
f.landmarks = raw_kps[i];
|
||||||
|
faces.push_back(f);
|
||||||
|
}
|
||||||
|
return faces;
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
struct TrtDeleter { template<class T> void operator()(T* p) const { delete p; } };
|
||||||
|
std::unique_ptr<nvinfer1::IRuntime, TrtDeleter> runtime_;
|
||||||
|
std::unique_ptr<nvinfer1::ICudaEngine, TrtDeleter> engine_;
|
||||||
|
std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
|
||||||
|
|
||||||
|
float conf_threshold_;
|
||||||
|
float nms_threshold_;
|
||||||
|
int fmc_{3};
|
||||||
|
|
||||||
|
std::string input_name_;
|
||||||
|
std::vector<std::string> output_names_;
|
||||||
|
void* d_input_ = nullptr;
|
||||||
|
std::vector<void*> d_outputs_;
|
||||||
|
mutable std::vector<std::vector<float>> host_outputs_;
|
||||||
|
std::vector<std::size_t> out_elem_counts_;
|
||||||
|
std::vector<int> out_last_dims_;
|
||||||
|
|
||||||
|
cudaStream_t stream_ = nullptr;
|
||||||
|
mutable std::mutex mu_;
|
||||||
|
};
|
||||||
+124
@@ -0,0 +1,124 @@
|
|||||||
|
#pragma once
|
||||||
|
#include <array>
|
||||||
|
#include <cstdint>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include <opencv2/core.hpp>
|
||||||
|
|
||||||
|
// ── Embedding ─────────────────────────────────────────────────────────────────
|
||||||
|
// 512-dim L2-normalised ArcFace embedding
|
||||||
|
using Embedding = std::array<float, 512>;
|
||||||
|
|
||||||
|
inline float cosine_similarity(const Embedding& a, const Embedding& b) {
|
||||||
|
float dot = 0.f;
|
||||||
|
for (int i = 0; i < 512; ++i) dot += a[i] * b[i];
|
||||||
|
return dot;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Frame ─────────────────────────────────────────────────────────────────────
|
||||||
|
// Raw sampled frame from the movie. eof=true is the pipeline shutdown sentinel:
|
||||||
|
// every node must forward it immediately without processing.
|
||||||
|
struct Frame {
|
||||||
|
cv::Mat image;
|
||||||
|
double timestamp_sec{0.0};
|
||||||
|
int64_t frame_idx{-1};
|
||||||
|
bool eof{false};
|
||||||
|
bool is_cut{false}; // true when a hard scene cut was detected before this frame
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── ArcFace alignment ─────────────────────────────────────────────────────────
|
||||||
|
// Canonical 5-point target positions for a 112×112 ArcFace crop.
|
||||||
|
// Landmark order: right-eye, left-eye, nose, right-mouth, left-mouth
|
||||||
|
// (matches SCRFD output order — no reordering needed).
|
||||||
|
inline constexpr float kArcFaceRef[5][2] = {
|
||||||
|
{38.2946f, 51.6963f},
|
||||||
|
{73.5318f, 51.5014f},
|
||||||
|
{56.0252f, 71.7366f},
|
||||||
|
{41.5493f, 92.3655f},
|
||||||
|
{70.7299f, 92.2041f},
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── DetectedFace ──────────────────────────────────────────────────────────────
|
||||||
|
// One face found by SCRFD in a Frame.
|
||||||
|
// Landmark order matches ArcFace convention (same as SCRFD output order):
|
||||||
|
// [0] right-eye-centre [1] left-eye-centre [2] nose
|
||||||
|
// [3] right-mouth [4] left-mouth
|
||||||
|
struct DetectedFace {
|
||||||
|
cv::Rect2f bbox;
|
||||||
|
std::array<cv::Point2f, 5> landmarks;
|
||||||
|
float confidence{0.f};
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── Pipeline messages ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
struct SceneFrame {
|
||||||
|
Frame source;
|
||||||
|
std::vector<DetectedFace> faces; // empty when no faces detected (or eof)
|
||||||
|
};
|
||||||
|
|
||||||
|
struct AlignedSceneFrame {
|
||||||
|
Frame source;
|
||||||
|
std::vector<DetectedFace> faces;
|
||||||
|
std::vector<cv::Mat> crops; // 112×112 BGR, ArcFace-ready; parallel to faces
|
||||||
|
};
|
||||||
|
|
||||||
|
struct EmbeddedSceneFrame {
|
||||||
|
Frame source;
|
||||||
|
std::vector<DetectedFace> faces;
|
||||||
|
std::vector<cv::Mat> crops; // forwarded for debug rendering downstream
|
||||||
|
std::vector<Embedding> embeddings;
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── Face tracking ─────────────────────────────────────────────────────────────
|
||||||
|
// Output of FaceTrackerFunc — EmbeddedSceneFrame augmented with per-detection
|
||||||
|
// track context. track_embeddings[i] is the L2-normalised running mean across
|
||||||
|
// the track's history; use it for identity matching when track_mature[i] is true.
|
||||||
|
|
||||||
|
struct TrackedSceneFrame {
|
||||||
|
Frame source;
|
||||||
|
std::vector<DetectedFace> faces;
|
||||||
|
std::vector<cv::Mat> crops;
|
||||||
|
std::vector<int> track_ids; // -1 = brand-new track this frame
|
||||||
|
std::vector<Embedding> embeddings; // per-frame raw (from embedder)
|
||||||
|
std::vector<Embedding> track_embeddings; // accumulated mean per track
|
||||||
|
std::vector<bool> track_mature; // true once track has ≥ min_frames obs.
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── Identity matching ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
struct IdentifiedActor {
|
||||||
|
int actor_idx{-1}; // index into ActorGallery::actors; -1 = unknown
|
||||||
|
int track_id{-1}; // face track ID from FaceTrackerFunc
|
||||||
|
std::string name;
|
||||||
|
std::string imdb_id;
|
||||||
|
float similarity{0.f}; // calibrated P(match) or cosine similarity; 0 for unknowns
|
||||||
|
cv::Rect2f bbox;
|
||||||
|
cv::Mat crop; // 112×112 aligned crop (stored as shared_ptr by KPN)
|
||||||
|
};
|
||||||
|
|
||||||
|
struct MatchedSceneFrame {
|
||||||
|
Frame source;
|
||||||
|
std::vector<IdentifiedActor> actors; // includes unknowns (actor_idx == -1)
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── Scene annotation ──────────────────────────────────────────────────────────
|
||||||
|
// Output of the scene tracker: one per sampled frame.
|
||||||
|
// visible_actors contains all actors still within their extinction window.
|
||||||
|
struct SceneAnnotation {
|
||||||
|
double timestamp_sec{0.0};
|
||||||
|
std::vector<IdentifiedActor> visible_actors;
|
||||||
|
bool eof{false};
|
||||||
|
};
|
||||||
|
|
||||||
|
// ── Actor gallery ─────────────────────────────────────────────────────────────
|
||||||
|
// Loaded once at startup; baked into the identity matcher.
|
||||||
|
struct ActorGallery {
|
||||||
|
struct Actor {
|
||||||
|
std::string imdb_id;
|
||||||
|
std::string name;
|
||||||
|
std::vector<Embedding> embeddings; // one per reference image
|
||||||
|
std::vector<std::string> source_images;
|
||||||
|
};
|
||||||
|
std::vector<Actor> actors;
|
||||||
|
};
|
||||||
Reference in New Issue
Block a user