forked from Keelrleek/SoulStream
KeelrLeek
a667b4b28e
The previous .git directory was found gutted (only info/ remained), so history is unrecoverable locally. This commit restores version control with the tree exactly as it stands: the working app, the preview-tracks launcher on the artist hero, the scrollable-track-previews design bundle (docs + approved mockups), and the Phase 0 validation artifacts.
469 lines
19 KiB
Rust
469 lines
19 KiB
Rust
//! Lightweight, local subject analysis for artist-hero composition.
|
|
//!
|
|
//! Person boxes come from the bundled YOLOX-Nano model through the Windows
|
|
//! ML runtime already present in Windows 10/11. Faces come from the native
|
|
//! Windows face detector, so no second model or inference runtime is shipped.
|
|
//! All failures are recoverable: callers fall back to a deterministic crop.
|
|
|
|
use crate::hero_crop::{FaceBox, SubjectBox};
|
|
|
|
#[cfg(windows)]
|
|
mod windows_detector {
|
|
use std::sync::OnceLock;
|
|
|
|
use image::{DynamicImage, GenericImageView as _, ImageEncoder as _};
|
|
use tauri::Manager as _;
|
|
use windows::{
|
|
core::{Interface as _, HSTRING},
|
|
Graphics::Imaging::{BitmapDecoder, BitmapPixelFormat, BitmapSize, SoftwareBitmap},
|
|
Media::FaceAnalysis::FaceDetector,
|
|
Storage::Streams::{DataWriter, InMemoryRandomAccessStream},
|
|
Win32::System::WinRT::{RoInitialize, RoUninitialize, RO_INIT_MULTITHREADED},
|
|
AI::MachineLearning::{
|
|
LearningModel, LearningModelBinding, LearningModelDevice, LearningModelDeviceKind,
|
|
LearningModelSession, TensorFloat,
|
|
},
|
|
};
|
|
|
|
use super::{FaceBox, SubjectBox};
|
|
|
|
const YOLOX_EDGE: u32 = 416;
|
|
const YOLOX_CANDIDATES: usize = 3549;
|
|
const YOLOX_VALUES: usize = 85;
|
|
const PERSON_SCORE_MIN: f32 = 0.25;
|
|
const PERSON_NMS_IOU: f32 = 0.45;
|
|
const FACE_ANALYSIS_MAX_EDGE: u32 = 960;
|
|
|
|
static PERSON_SESSION: OnceLock<Result<LearningModelSession, String>> = OnceLock::new();
|
|
static FACE_DETECTOR: OnceLock<Result<FaceDetector, String>> = OnceLock::new();
|
|
|
|
struct WinRtGuard;
|
|
|
|
impl Drop for WinRtGuard {
|
|
fn drop(&mut self) {
|
|
unsafe { RoUninitialize() };
|
|
}
|
|
}
|
|
|
|
fn initialize_winrt() -> Result<Option<WinRtGuard>, String> {
|
|
match unsafe { RoInitialize(RO_INIT_MULTITHREADED) } {
|
|
Ok(()) => Ok(Some(WinRtGuard)),
|
|
Err(error) if error.code().0 as u32 == 0x8001_0106 => {
|
|
// The current thread already has a different apartment. WinRT
|
|
// remains usable there; it simply must not be uninitialized by us.
|
|
Ok(None)
|
|
}
|
|
Err(error) => Err(format!("Windows Runtime initialization failed: {error}")),
|
|
}
|
|
}
|
|
|
|
fn model_path(app: &tauri::AppHandle) -> Result<std::path::PathBuf, String> {
|
|
let bundled = app
|
|
.path()
|
|
.resolve(
|
|
"resources/yolox_nano.onnx",
|
|
tauri::path::BaseDirectory::Resource,
|
|
)
|
|
.map_err(|error| error.to_string())?;
|
|
if bundled.is_file() {
|
|
return Ok(bundled);
|
|
}
|
|
let development = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
|
.join("resources")
|
|
.join("yolox_nano.onnx");
|
|
if development.is_file() {
|
|
return Ok(development);
|
|
}
|
|
Err("bundled YOLOX-Nano model is missing".into())
|
|
}
|
|
|
|
fn person_session_from_path(path: &std::path::Path) -> Result<LearningModelSession, String> {
|
|
PERSON_SESSION
|
|
.get_or_init(|| {
|
|
let model = LearningModel::LoadFromFilePath(&HSTRING::from(
|
|
path.to_string_lossy().as_ref(),
|
|
))
|
|
.map_err(|error| format!("could not load YOLOX-Nano: {error}"))?;
|
|
let device = LearningModelDevice::Create(LearningModelDeviceKind::Cpu)
|
|
.map_err(|error| format!("could not create WinML CPU device: {error}"))?;
|
|
LearningModelSession::CreateFromModelOnDevice(&model, &device)
|
|
.map_err(|error| format!("could not create YOLOX session: {error}"))
|
|
})
|
|
.clone()
|
|
}
|
|
|
|
fn person_session(app: &tauri::AppHandle) -> Result<LearningModelSession, String> {
|
|
person_session_from_path(&model_path(app)?)
|
|
}
|
|
|
|
fn face_detector() -> Result<FaceDetector, String> {
|
|
FACE_DETECTOR
|
|
.get_or_init(|| {
|
|
if !FaceDetector::IsSupported()
|
|
.map_err(|error| format!("could not query face detector: {error}"))?
|
|
{
|
|
return Err("Windows face detection is unavailable".into());
|
|
}
|
|
let detector = FaceDetector::CreateAsync()
|
|
.map_err(|error| format!("could not create face detector: {error}"))?
|
|
.get()
|
|
.map_err(|error| format!("could not initialize face detector: {error}"))?;
|
|
// Small faces matter in group press photos. Analysis is capped
|
|
// at 960 px, so 18 px still avoids tiny texture false positives.
|
|
let _ = detector.SetMinDetectableFaceSize(BitmapSize {
|
|
Width: 18,
|
|
Height: 18,
|
|
});
|
|
Ok(detector)
|
|
})
|
|
.clone()
|
|
}
|
|
|
|
pub(super) fn detect(
|
|
app: &tauri::AppHandle,
|
|
image: &DynamicImage,
|
|
) -> Result<(Vec<SubjectBox>, Vec<FaceBox>), String> {
|
|
let _winrt = initialize_winrt()?;
|
|
let subjects = detect_people(app, image)?;
|
|
// Face refinement is optional. A native-detector failure must not throw
|
|
// away valid person boxes from WinML.
|
|
let faces = match detect_faces(image) {
|
|
Ok(faces) => faces,
|
|
Err(error) => {
|
|
log::warn!("artist hero face refinement unavailable: {error}");
|
|
Vec::new()
|
|
}
|
|
};
|
|
Ok((subjects, faces))
|
|
}
|
|
|
|
pub(super) fn warm_up(app: tauri::AppHandle) {
|
|
std::thread::Builder::new()
|
|
.name("hero-vision-warmup".into())
|
|
.spawn(move || {
|
|
let Ok(_winrt) = initialize_winrt() else {
|
|
return;
|
|
};
|
|
if let Err(error) = person_session(&app) {
|
|
log::warn!("artist hero person detector warm-up failed: {error}");
|
|
}
|
|
if let Err(error) = face_detector() {
|
|
log::warn!("artist hero face detector warm-up failed: {error}");
|
|
}
|
|
})
|
|
.ok();
|
|
}
|
|
|
|
fn detect_people(
|
|
app: &tauri::AppHandle,
|
|
image: &DynamicImage,
|
|
) -> Result<Vec<SubjectBox>, String> {
|
|
detect_people_with_model_path(&model_path(app)?, image)
|
|
}
|
|
|
|
fn detect_people_with_model_path(
|
|
model_path: &std::path::Path,
|
|
image: &DynamicImage,
|
|
) -> Result<Vec<SubjectBox>, String> {
|
|
let (source_width, source_height) = image.dimensions();
|
|
if source_width == 0 || source_height == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let scale =
|
|
(YOLOX_EDGE as f32 / source_width as f32).min(YOLOX_EDGE as f32 / source_height as f32);
|
|
let resized_width = ((source_width as f32 * scale).floor() as u32).clamp(1, YOLOX_EDGE);
|
|
let resized_height = ((source_height as f32 * scale).floor() as u32).clamp(1, YOLOX_EDGE);
|
|
let rgb = image.to_rgb8();
|
|
let resized = image::imageops::resize(
|
|
&rgb,
|
|
resized_width,
|
|
resized_height,
|
|
image::imageops::FilterType::Triangle,
|
|
);
|
|
|
|
let plane_len = (YOLOX_EDGE * YOLOX_EDGE) as usize;
|
|
let mut tensor_data = vec![114.0f32; plane_len * 3];
|
|
for y in 0..resized_height {
|
|
for x in 0..resized_width {
|
|
let pixel = resized.get_pixel(x, y).0;
|
|
let offset = (y * YOLOX_EDGE + x) as usize;
|
|
// Official YOLOX preprocessing consumes OpenCV BGR without
|
|
// normalization, transposed directly into NCHW planes.
|
|
tensor_data[offset] = f32::from(pixel[2]);
|
|
tensor_data[plane_len + offset] = f32::from(pixel[1]);
|
|
tensor_data[plane_len * 2 + offset] = f32::from(pixel[0]);
|
|
}
|
|
}
|
|
|
|
let input = TensorFloat::CreateFromShapeArrayAndDataArray(
|
|
&[1, 3, i64::from(YOLOX_EDGE), i64::from(YOLOX_EDGE)],
|
|
&tensor_data,
|
|
)
|
|
.map_err(|error| format!("could not create YOLOX input: {error}"))?;
|
|
let session = person_session_from_path(model_path)?;
|
|
let binding = LearningModelBinding::CreateFromSession(&session)
|
|
.map_err(|error| format!("could not create YOLOX binding: {error}"))?;
|
|
binding
|
|
.Bind(&HSTRING::from("images"), &input)
|
|
.map_err(|error| format!("could not bind YOLOX input: {error}"))?;
|
|
let result = session
|
|
.Evaluate(&binding, &HSTRING::from("artist-hero"))
|
|
.map_err(|error| format!("YOLOX inference failed: {error}"))?;
|
|
if !result
|
|
.Succeeded()
|
|
.map_err(|error| format!("could not read YOLOX status: {error}"))?
|
|
{
|
|
return Err(format!(
|
|
"YOLOX evaluation failed with status {}",
|
|
result.ErrorStatus().unwrap_or_default()
|
|
));
|
|
}
|
|
let output = result
|
|
.Outputs()
|
|
.and_then(|outputs| outputs.Lookup(&HSTRING::from("output")))
|
|
.map_err(|error| format!("YOLOX output is missing: {error}"))?;
|
|
let tensor: TensorFloat = output
|
|
.cast()
|
|
.map_err(|error| format!("YOLOX output has the wrong type: {error}"))?;
|
|
let view = tensor
|
|
.GetAsVectorView()
|
|
.map_err(|error| format!("could not read YOLOX output: {error}"))?;
|
|
let values: Vec<f32> = view.into_iter().collect();
|
|
if values.len() != YOLOX_CANDIDATES * YOLOX_VALUES {
|
|
return Err(format!("unexpected YOLOX output length: {}", values.len()));
|
|
}
|
|
|
|
let mut candidates = Vec::new();
|
|
let mut candidate_index = 0usize;
|
|
for stride in [8u32, 16, 32] {
|
|
let grid = YOLOX_EDGE / stride;
|
|
for gy in 0..grid {
|
|
for gx in 0..grid {
|
|
let row = &values
|
|
[candidate_index * YOLOX_VALUES..(candidate_index + 1) * YOLOX_VALUES];
|
|
candidate_index += 1;
|
|
let confidence = row[4] * row[5];
|
|
if !confidence.is_finite() || confidence < PERSON_SCORE_MIN {
|
|
continue;
|
|
}
|
|
let center_x = (row[0] + gx as f32) * stride as f32 / scale;
|
|
let center_y = (row[1] + gy as f32) * stride as f32 / scale;
|
|
let width = row[2].clamp(-10.0, 10.0).exp() * stride as f32 / scale;
|
|
let height = row[3].clamp(-10.0, 10.0).exp() * stride as f32 / scale;
|
|
if !center_x.is_finite()
|
|
|| !center_y.is_finite()
|
|
|| !width.is_finite()
|
|
|| !height.is_finite()
|
|
{
|
|
continue;
|
|
}
|
|
let left = (center_x - width * 0.5).clamp(0.0, source_width as f32);
|
|
let top = (center_y - height * 0.5).clamp(0.0, source_height as f32);
|
|
let right = (center_x + width * 0.5).clamp(0.0, source_width as f32);
|
|
let bottom = (center_y + height * 0.5).clamp(0.0, source_height as f32);
|
|
if right - left < 8.0 || bottom - top < 16.0 {
|
|
continue;
|
|
}
|
|
candidates.push(SubjectBox {
|
|
x: left,
|
|
y: top,
|
|
width: right - left,
|
|
height: bottom - top,
|
|
confidence,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
debug_assert_eq!(candidate_index, YOLOX_CANDIDATES);
|
|
Ok(non_maximum_suppression(candidates, PERSON_NMS_IOU, 24))
|
|
}
|
|
|
|
fn detect_faces(image: &DynamicImage) -> Result<Vec<FaceBox>, String> {
|
|
let (source_width, source_height) = image.dimensions();
|
|
if source_width == 0 || source_height == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
let detection = image.thumbnail(FACE_ANALYSIS_MAX_EDGE, FACE_ANALYSIS_MAX_EDGE);
|
|
let gray = detection.to_luma8();
|
|
let (width, height) = gray.dimensions();
|
|
let mut png = Vec::new();
|
|
image::codecs::png::PngEncoder::new(&mut png)
|
|
.write_image(gray.as_raw(), width, height, image::ExtendedColorType::L8)
|
|
.map_err(|error| format!("could not encode face-analysis raster: {error}"))?;
|
|
|
|
let stream = InMemoryRandomAccessStream::new()
|
|
.map_err(|error| format!("could not allocate face-analysis stream: {error}"))?;
|
|
let writer = DataWriter::CreateDataWriter(&stream)
|
|
.map_err(|error| format!("could not create face-analysis writer: {error}"))?;
|
|
writer
|
|
.WriteBytes(&png)
|
|
.map_err(|error| format!("could not write face-analysis raster: {error}"))?;
|
|
writer
|
|
.StoreAsync()
|
|
.and_then(|operation| operation.get())
|
|
.map_err(|error| format!("could not store face-analysis raster: {error}"))?;
|
|
writer
|
|
.FlushAsync()
|
|
.and_then(|operation| operation.get())
|
|
.map_err(|error| format!("could not flush face-analysis raster: {error}"))?;
|
|
let _ = writer.DetachStream();
|
|
stream
|
|
.Seek(0)
|
|
.map_err(|error| format!("could not rewind face-analysis raster: {error}"))?;
|
|
let decoder = BitmapDecoder::CreateAsync(&stream)
|
|
.and_then(|operation| operation.get())
|
|
.map_err(|error| format!("could not decode face-analysis raster: {error}"))?;
|
|
let decoded_bitmap = decoder
|
|
.GetSoftwareBitmapAsync()
|
|
.and_then(|operation| operation.get())
|
|
.map_err(|error| format!("could not prepare face-analysis bitmap: {error}"))?;
|
|
let bitmap = SoftwareBitmap::Convert(&decoded_bitmap, BitmapPixelFormat::Gray8)
|
|
.map_err(|error| format!("could not convert face-analysis bitmap: {error}"))?;
|
|
let found = face_detector()?
|
|
.DetectFacesAsync(&bitmap)
|
|
.and_then(|operation| operation.get())
|
|
.map_err(|error| format!("face detection failed: {error}"))?;
|
|
|
|
let scale_x = source_width as f32 / width as f32;
|
|
let scale_y = source_height as f32 / height as f32;
|
|
let mut faces = Vec::new();
|
|
for detected in &found {
|
|
let bounds = detected
|
|
.FaceBox()
|
|
.map_err(|error| format!("could not read face bounds: {error}"))?;
|
|
faces.push(FaceBox {
|
|
x: bounds.X as f32 * scale_x,
|
|
y: bounds.Y as f32 * scale_y,
|
|
width: bounds.Width as f32 * scale_x,
|
|
height: bounds.Height as f32 * scale_y,
|
|
confidence: 1.0,
|
|
});
|
|
}
|
|
Ok(faces)
|
|
}
|
|
|
|
fn non_maximum_suppression(
|
|
mut candidates: Vec<SubjectBox>,
|
|
threshold: f32,
|
|
limit: usize,
|
|
) -> Vec<SubjectBox> {
|
|
candidates.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
|
|
let mut kept = Vec::new();
|
|
'candidate: for candidate in candidates {
|
|
for accepted in &kept {
|
|
if intersection_over_union(candidate, *accepted) > threshold {
|
|
continue 'candidate;
|
|
}
|
|
}
|
|
kept.push(candidate);
|
|
if kept.len() == limit {
|
|
break;
|
|
}
|
|
}
|
|
kept
|
|
}
|
|
|
|
fn intersection_over_union(a: SubjectBox, b: SubjectBox) -> f32 {
|
|
let left = a.x.max(b.x);
|
|
let top = a.y.max(b.y);
|
|
let right = (a.x + a.width).min(b.x + b.width);
|
|
let bottom = (a.y + a.height).min(b.y + b.height);
|
|
let intersection = (right - left).max(0.0) * (bottom - top).max(0.0);
|
|
let union = a.width * a.height + b.width * b.height - intersection;
|
|
if union > 0.0 {
|
|
intersection / union
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn nms_keeps_the_best_overlapping_person() {
|
|
let kept = non_maximum_suppression(
|
|
vec![
|
|
SubjectBox {
|
|
x: 0.0,
|
|
y: 0.0,
|
|
width: 100.0,
|
|
height: 200.0,
|
|
confidence: 0.9,
|
|
},
|
|
SubjectBox {
|
|
x: 2.0,
|
|
y: 4.0,
|
|
width: 100.0,
|
|
height: 200.0,
|
|
confidence: 0.8,
|
|
},
|
|
SubjectBox {
|
|
x: 300.0,
|
|
y: 0.0,
|
|
width: 100.0,
|
|
height: 200.0,
|
|
confidence: 0.7,
|
|
},
|
|
],
|
|
0.45,
|
|
24,
|
|
);
|
|
assert_eq!(kept.len(), 2);
|
|
assert_eq!(kept[0].confidence, 0.9);
|
|
}
|
|
|
|
/// Manual, reusable integration probe. It is ignored in normal CI;
|
|
/// provide an image through SOULSTREAM_HERO_FIXTURE when validating a
|
|
/// new detector/model or an especially difficult artist photo.
|
|
#[test]
|
|
#[ignore]
|
|
fn fixture_detection_and_crop_probe() {
|
|
let fixture = std::env::var("SOULSTREAM_HERO_FIXTURE")
|
|
.expect("set SOULSTREAM_HERO_FIXTURE to a source image");
|
|
let image = image::open(fixture).expect("fixture must decode");
|
|
let _winrt = initialize_winrt().expect("WinRT must initialize");
|
|
let model = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
|
.join("resources")
|
|
.join("yolox_nano.onnx");
|
|
let subjects = detect_people_with_model_path(&model, &image)
|
|
.expect("person inference must succeed");
|
|
let faces = detect_faces(&image).expect("face detection must succeed");
|
|
let plan =
|
|
crate::hero_crop::plan_hero_crop(image.width(), image.height(), &subjects, &faces)
|
|
.expect("fixture must plan");
|
|
eprintln!("subjects={subjects:#?}");
|
|
eprintln!("faces={faces:#?}");
|
|
eprintln!("plan={plan:#?}");
|
|
assert!(plan.crop.width > 0 && plan.crop.height > 0);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Detects people and faces in source-image coordinates. This is deliberately
|
|
/// a fallible enhancement: hero processing must continue when vision APIs are
|
|
/// unavailable (including non-Windows development builds).
|
|
pub(crate) fn detect_subjects(
|
|
app: &tauri::AppHandle,
|
|
image: &image::DynamicImage,
|
|
) -> Result<(Vec<SubjectBox>, Vec<FaceBox>), String> {
|
|
#[cfg(windows)]
|
|
{
|
|
windows_detector::detect(app, image)
|
|
}
|
|
#[cfg(not(windows))]
|
|
{
|
|
let _ = (app, image);
|
|
Err("content-aware hero detection is only available on Windows".into())
|
|
}
|
|
}
|
|
|
|
pub(crate) fn warm_up(app: tauri::AppHandle) {
|
|
#[cfg(windows)]
|
|
windows_detector::warm_up(app);
|
|
#[cfg(not(windows))]
|
|
let _ = app;
|
|
}
|