mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-14 19:27:01 -04:00
Brandon Thomas committed
commit
cbf9935b6e
31 files changed
+2227
-16
No files matched your search
Generated
+13
@@ -2882,6 +2882,7 @@ dependencies = [
|
||||
"lightcraft-codecs",
|
||||
"lightcraft-color",
|
||||
"lightcraft-develop",
|
||||
"lightcraft-faces",
|
||||
"lightcraft-geom",
|
||||
"lightcraft-gpu",
|
||||
"lightcraft-merge",
|
||||
@@ -2903,6 +2904,17 @@ dependencies = [
|
||||
"web-time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lightcraft-faces"
|
||||
version = "0.4.0"
|
||||
dependencies = [
|
||||
"proptest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"thiserror 2.0.21",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lightcraft-fetch"
|
||||
version = "0.4.0"
|
||||
@@ -3084,6 +3096,7 @@ dependencies = [
|
||||
"lightcraft-color",
|
||||
"lightcraft-develop",
|
||||
"lightcraft-engine",
|
||||
"lightcraft-faces",
|
||||
"lightcraft-geom",
|
||||
"lightcraft-meta",
|
||||
"lightcraft-pipeline",
|
||||
|
||||
@@ -12,6 +12,7 @@ repository = "https://github.com/storytold/lightcraft"
|
||||
[workspace.dependencies]
|
||||
# internal
|
||||
lightcraft-geom = { path = "crates/geom" }
|
||||
lightcraft-faces = { path = "crates/faces" }
|
||||
lightcraft-sysmem = { path = "crates/sysmem" }
|
||||
lightcraft-color = { path = "crates/color" }
|
||||
lightcraft-raster = { path = "crates/raster" }
|
||||
@@ -34,6 +35,7 @@ lightcraft-segment = { path = "crates/segment" }
|
||||
|
||||
# foundation
|
||||
thiserror = "2"
|
||||
sha2 = "0.10"
|
||||
serde = { version = "1", features = ["derive", "rc"] }
|
||||
serde_json = { version = "1", features = ["float_roundtrip"] }
|
||||
log = "0.4"
|
||||
|
||||
@@ -732,12 +732,12 @@ fn snapshot(args: &[String]) -> Result<(), String> {
|
||||
let t0 = Instant::now();
|
||||
let mut session = match &library {
|
||||
Some(dir) => {
|
||||
let mut s = Session::new().with_fs();
|
||||
let mut s = Session::new().with_fs().with_default_face_models();
|
||||
s.open_library(dir, false).map_err(|e| library_error(dir, e))?;
|
||||
s
|
||||
}
|
||||
None if files.is_empty() => Session::with_demo().with_fs(),
|
||||
None => Session::new().with_fs(),
|
||||
None if files.is_empty() => Session::with_demo().with_fs().with_default_face_models(),
|
||||
None => Session::new().with_fs().with_default_face_models(),
|
||||
};
|
||||
// Apply compute policy before imports or the first scripted query can discover an adapter.
|
||||
session.execute("app.gpu", &json!({"enabled": false})).map_err(|e| e.to_string())?;
|
||||
|
||||
@@ -209,6 +209,8 @@ fn services(originals: Originals, backend: Option<Backend>, files: Files, frozen
|
||||
pick_preset_files: None,
|
||||
pick_tracklog: None,
|
||||
pick_lightroom_catalog: None,
|
||||
// Face models need a folder to live in; the web has none.
|
||||
pick_model_file: None,
|
||||
save_preset_file: None,
|
||||
pick_curve_preset_files: None,
|
||||
save_curve_preset_file: None,
|
||||
|
||||
@@ -400,6 +400,14 @@ fn services(ctx: egui::Context, log_file: Option<&std::path::Path>) -> Services
|
||||
.map(|p| p.to_string_lossy().to_string())
|
||||
.collect()
|
||||
})),
|
||||
pick_model_file: Some(Box::new(|| {
|
||||
rfd::FileDialog::new()
|
||||
.set_title(lightcraft_ui_egui::i18n::tr("Add a Face Recognition Model"))
|
||||
.add_filter_nocase(lightcraft_ui_egui::i18n::tr("Face models (ONNX)"), &["onnx"])
|
||||
.pick_file()
|
||||
.map(|p| vec![p.to_string_lossy().to_string()])
|
||||
.unwrap_or_default()
|
||||
})),
|
||||
pick_tracklog: Some(Box::new(|| {
|
||||
rfd::FileDialog::new()
|
||||
.set_title(lightcraft_ui_egui::i18n::tr("Auto-Tag from Tracklog"))
|
||||
@@ -446,12 +454,19 @@ fn services(ctx: egui::Context, log_file: Option<&std::path::Path>) -> Services
|
||||
}
|
||||
}
|
||||
|
||||
/// The library session with the folder face models are kept in (`<config>/models`).
|
||||
fn open_session(in_memory: bool, dir: Option<std::path::PathBuf>, seed_demo: bool) -> (Session, Option<LibraryProblem>) {
|
||||
let (mut s, problem) = open_library_session(in_memory, dir, seed_demo);
|
||||
s.face_models_dir = lightcraft_engine::config::default_face_models_dir();
|
||||
(s, problem)
|
||||
}
|
||||
|
||||
/// The persistent library session, or with `--memory` an in-memory one (demo photos).
|
||||
///
|
||||
/// If the library can't be opened the session is empty and in memory — never seeded with the
|
||||
/// demo photos, never written anywhere — and the problem is returned: the window then says so
|
||||
/// and offers Try Again / Choose Another Library… / Continue Without Saving / Quit (issue #100).
|
||||
fn open_session(in_memory: bool, dir: Option<std::path::PathBuf>, seed_demo: bool) -> (Session, Option<LibraryProblem>) {
|
||||
fn open_library_session(in_memory: bool, dir: Option<std::path::PathBuf>, seed_demo: bool) -> (Session, Option<LibraryProblem>) {
|
||||
if in_memory {
|
||||
return (if seed_demo { Session::with_demo() } else { Session::new() }.with_fs().with_system_clock(), None);
|
||||
}
|
||||
|
||||
@@ -22,6 +22,7 @@ lightcraft-preview = { workspace = true }
|
||||
lightcraft-codecs = { workspace = true }
|
||||
lightcraft-raw = { workspace = true }
|
||||
lightcraft-meta = { workspace = true }
|
||||
lightcraft-faces = { workspace = true }
|
||||
miniz_oxide = { workspace = true }
|
||||
lightcraft-merge = { workspace = true }
|
||||
rayon = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
//! Face model commands (`faces.*`): which models exist, inspecting a file the user dropped in, installing
|
||||
//! it once its licence is accepted, removing it, and the faces on/off setting.
|
||||
//!
|
||||
//! Models live in the folder the host provides (`Session::face_models_dir`): one subfolder per model with
|
||||
//! `model.onnx`, `face-model.json` (its manifest) and `installed.json` (what the user accepted), plus
|
||||
//! `settings.json` beside them. Model files and manifests are hostile input: sizes are capped, ids are
|
||||
//! validated before they become folder names, and nothing outside a model's own folder is ever touched.
|
||||
|
||||
use std::io::Read;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use lightcraft_faces::hash::sha256_file;
|
||||
use lightcraft_faces::manifest::{self, MAX_MANIFEST_BYTES, MAX_MODEL_BYTES, ModelManifest, Role};
|
||||
use lightcraft_faces::suggest::{Suggestion, suggest};
|
||||
use lightcraft_faces::{known, onnx};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{CommandSpec, always, bad, cmd, str_param};
|
||||
use crate::{EngineError, Result, Session};
|
||||
|
||||
#[derive(Default, Serialize, Deserialize)]
|
||||
#[serde(default, rename_all = "camelCase")]
|
||||
struct FaceSettings {
|
||||
enabled: bool,
|
||||
/// The installed recogniser in use.
|
||||
embedder: Option<String>,
|
||||
}
|
||||
|
||||
fn fail(what: &str, e: impl std::fmt::Display) -> EngineError {
|
||||
EngineError::Other(format!("{what}: {e}"))
|
||||
}
|
||||
|
||||
fn models_dir(s: &Session, cmd: &str) -> Result<PathBuf> {
|
||||
s.face_models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep face models (the desktop app does)"))
|
||||
}
|
||||
|
||||
/// Read at most `max` bytes of a small file.
|
||||
fn read_capped(path: &Path, max: usize) -> Option<Vec<u8>> {
|
||||
let mut buf = Vec::new();
|
||||
std::fs::File::open(path).ok()?.take(max as u64 + 1).read_to_end(&mut buf).ok()?;
|
||||
(buf.len() <= max).then_some(buf)
|
||||
}
|
||||
|
||||
fn read_settings(dir: &Path) -> FaceSettings {
|
||||
read_capped(&dir.join("settings.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> {
|
||||
let part = path.with_extension("part");
|
||||
std::fs::write(&part, bytes).map_err(|e| fail(&format!("could not write {}", path.display()), e))?;
|
||||
std::fs::rename(&part, path).map_err(|e| {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
fail(&format!("could not write {}", path.display()), e)
|
||||
})
|
||||
}
|
||||
|
||||
fn write_settings(dir: &Path, st: &FaceSettings) -> Result<()> {
|
||||
std::fs::create_dir_all(dir).map_err(|e| fail("could not create the face models folder", e))?;
|
||||
write_atomic(&dir.join("settings.json"), &serde_json::to_vec_pretty(st).map_err(|e| fail("settings", e))?)
|
||||
}
|
||||
|
||||
struct Installed {
|
||||
manifest: ModelManifest,
|
||||
accepted: Value,
|
||||
}
|
||||
|
||||
/// The models in the folder. A folder that is not a valid model (wrong name, no model file, a manifest that
|
||||
/// does not validate or does not match its folder) is ignored, never trusted.
|
||||
fn installed_models(dir: &Path) -> Vec<Installed> {
|
||||
let Ok(entries) = std::fs::read_dir(dir) else { return Vec::new() };
|
||||
let mut out = Vec::new();
|
||||
for e in entries.flatten() {
|
||||
let p = e.path();
|
||||
let Some(name) = p.file_name().and_then(|n| n.to_str()).map(str::to_string) else { continue };
|
||||
if !manifest::valid_id(&name) || !p.is_dir() || !p.join("model.onnx").is_file() {
|
||||
continue;
|
||||
}
|
||||
let Some(m) = read_capped(&p.join("face-model.json"), MAX_MANIFEST_BYTES).and_then(|b| manifest::parse(&b).ok()) else { continue };
|
||||
if m.id != name {
|
||||
continue;
|
||||
}
|
||||
let accepted = read_capped(&p.join("installed.json"), 4096).and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or(Value::Null);
|
||||
out.push(Installed { manifest: m, accepted });
|
||||
}
|
||||
out.sort_by(|a, b| a.manifest.id.cmp(&b.manifest.id));
|
||||
out
|
||||
}
|
||||
|
||||
fn row(m: &ModelManifest, installed: bool, selected: bool, accepted: &Value) -> Value {
|
||||
json!({
|
||||
"id": m.id,
|
||||
"name": m.name,
|
||||
"role": m.role,
|
||||
"licence": m.licence,
|
||||
"provenance": m.provenance,
|
||||
"source": m.source,
|
||||
"sizeBytes": m.size_bytes,
|
||||
"sha256": m.sha256,
|
||||
"known": known::all().iter().any(|k| k.id == m.id),
|
||||
"installed": installed,
|
||||
"selected": selected,
|
||||
"accepted": accepted,
|
||||
})
|
||||
}
|
||||
|
||||
struct Inspection {
|
||||
sha256: String,
|
||||
size: u64,
|
||||
file_name: String,
|
||||
suggestion: Suggestion,
|
||||
}
|
||||
|
||||
fn inspect_file(path: &str, cmd: &str) -> Result<Inspection> {
|
||||
let p = Path::new(path);
|
||||
let meta = std::fs::metadata(p).map_err(|e| bad(cmd, format!("cannot read `{path}`: {e}")))?;
|
||||
if !meta.is_file() {
|
||||
return Err(bad(cmd, format!("`{path}` is not a file")));
|
||||
}
|
||||
if meta.len() == 0 || meta.len() > MAX_MODEL_BYTES {
|
||||
return Err(bad(cmd, "a model file must be between 1 byte and 4 GiB"));
|
||||
}
|
||||
let sha256 = sha256_file(p).map_err(|e| fail(&format!("could not read {path}"), e))?;
|
||||
let file_name = p.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default();
|
||||
let suggestion = match onnx::probe_path(p) {
|
||||
Ok(info) => suggest(&info, &sha256, meta.len(), &file_name),
|
||||
Err(e) => Suggestion::Unsupported(e.to_string()),
|
||||
};
|
||||
Ok(Inspection { sha256, size: meta.len(), file_name, suggestion })
|
||||
}
|
||||
|
||||
fn list(s: &mut Session, _: &Value) -> Result<Value> {
|
||||
let dir = s.face_models_dir.clone();
|
||||
let (settings, installed) = match &dir {
|
||||
Some(d) => (read_settings(d), installed_models(d)),
|
||||
None => (FaceSettings::default(), Vec::new()),
|
||||
};
|
||||
let on_disk = |id: &str| installed.iter().find(|i| i.manifest.id == id);
|
||||
let mut models: Vec<Value> = Vec::new();
|
||||
for k in known::all() {
|
||||
let found = on_disk(&k.id);
|
||||
models.push(row(&k, found.is_some(), settings.embedder.as_deref() == Some(k.id.as_str()), found.map_or(&Value::Null, |f| &f.accepted)));
|
||||
}
|
||||
for i in &installed {
|
||||
if !known::all().iter().any(|k| k.id == i.manifest.id) {
|
||||
models.push(row(&i.manifest, true, settings.embedder.as_deref() == Some(i.manifest.id.as_str()), &i.accepted));
|
||||
}
|
||||
}
|
||||
Ok(json!({
|
||||
"dir": dir.as_ref().map(|d| d.display().to_string()),
|
||||
"enabled": settings.enabled,
|
||||
"embedder": settings.embedder,
|
||||
// running a model is not built in yet: models can be installed and chosen, not used
|
||||
"runtime": false,
|
||||
"models": models,
|
||||
}))
|
||||
}
|
||||
|
||||
fn inspect(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.inspect";
|
||||
let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?;
|
||||
let ins = inspect_file(path, C)?;
|
||||
let installed = s.face_models_dir.as_deref().map(installed_models).unwrap_or_default();
|
||||
let (kind, model, assumptions, reason) = match &ins.suggestion {
|
||||
Suggestion::Known(m) => ("known", Some(m), Vec::new(), None),
|
||||
Suggestion::Draft { manifest, assumptions } => ("draft", Some(manifest), assumptions.clone(), None),
|
||||
Suggestion::Unsupported(why) => ("unsupported", None, Vec::new(), Some(why.clone())),
|
||||
};
|
||||
let already = model.is_some_and(|m| installed.iter().any(|i| i.manifest.sha256 == m.sha256 && i.manifest.id == m.id));
|
||||
Ok(json!({
|
||||
"path": path,
|
||||
"fileName": ins.file_name,
|
||||
"sizeBytes": ins.size,
|
||||
"sha256": ins.sha256,
|
||||
"kind": kind,
|
||||
"model": model.map(|m| row(m, already, false, &Value::Null)),
|
||||
"assumptions": assumptions,
|
||||
"reason": reason,
|
||||
"alreadyInstalled": already,
|
||||
}))
|
||||
}
|
||||
|
||||
fn install(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.install";
|
||||
let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?;
|
||||
let acknowledged = p.get("acknowledged").and_then(Value::as_bool).unwrap_or(false);
|
||||
let dir = models_dir(s, C)?;
|
||||
let ins = inspect_file(path, C)?;
|
||||
let m = match ins.suggestion {
|
||||
Suggestion::Known(m) | Suggestion::Draft { manifest: m, .. } => m,
|
||||
Suggestion::Unsupported(why) => return Err(bad(C, format!("this model cannot be used yet: {why}"))),
|
||||
};
|
||||
if m.role != Role::Embedder {
|
||||
return Err(bad(C, "only face recognition models can be installed for now"));
|
||||
}
|
||||
manifest::validate(&m).map_err(|e| bad(C, e.to_string()))?;
|
||||
if !acknowledged {
|
||||
return Err(bad(C, "the licence has not been accepted: show the user its terms and pass `acknowledged: true` once they agree"));
|
||||
}
|
||||
let home = dir.join(&m.id);
|
||||
std::fs::create_dir_all(&home).map_err(|e| fail("could not create the model's folder", e))?;
|
||||
let (part, final_path) = (home.join("model.onnx.part"), home.join("model.onnx"));
|
||||
let copied =
|
||||
std::fs::copy(path, &part).map_err(|e| fail("could not copy the model (is the disk full?)", e)).and_then(|_| match sha256_file(&part) {
|
||||
Ok(h) if h == ins.sha256 => Ok(()),
|
||||
Ok(_) => Err(EngineError::Other("the copy does not match the original (the file changed while copying?)".into())),
|
||||
Err(e) => Err(fail("could not check the copy", e)),
|
||||
});
|
||||
if let Err(e) = copied.and_then(|_| std::fs::rename(&part, &final_path).map_err(|e| fail("could not finish the copy", e))) {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(e);
|
||||
}
|
||||
write_atomic(&home.join("face-model.json"), &serde_json::to_vec_pretty(&m).map_err(|e| fail("manifest", e))?)?;
|
||||
let accepted = json!({"acceptedAt": (s.clock)(), "licence": m.licence.name, "commercial": m.licence.commercial, "fileName": ins.file_name});
|
||||
write_atomic(&home.join("installed.json"), &serde_json::to_vec_pretty(&accepted).map_err(|e| fail("record", e))?)?;
|
||||
Ok(json!({"installed": row(&m, true, false, &accepted)}))
|
||||
}
|
||||
|
||||
fn remove(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.remove";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
if !manifest::valid_id(id) {
|
||||
return Err(bad(C, "not a model id"));
|
||||
}
|
||||
let home = dir.join(id);
|
||||
// only a folder that is a model (it has a manifest) is ever deleted
|
||||
if !home.join("face-model.json").is_file() {
|
||||
return Err(bad(C, "that model is not installed"));
|
||||
}
|
||||
std::fs::remove_dir_all(&home).map_err(|e| fail("could not remove the model", e))?;
|
||||
let mut st = read_settings(&dir);
|
||||
if st.embedder.as_deref() == Some(id) {
|
||||
st.embedder = None;
|
||||
write_settings(&dir, &st)?;
|
||||
}
|
||||
Ok(json!({"removed": id}))
|
||||
}
|
||||
|
||||
fn select(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.select";
|
||||
let dir = models_dir(s, C)?;
|
||||
let id = match p.get("id") {
|
||||
None | Some(Value::Null) => None,
|
||||
Some(v) => Some(v.as_str().ok_or_else(|| bad(C, "`id` must be a model id or null"))?.to_string()),
|
||||
};
|
||||
if let Some(id) = &id
|
||||
&& !installed_models(&dir).iter().any(|i| &i.manifest.id == id && i.manifest.role == Role::Embedder)
|
||||
{
|
||||
return Err(bad(C, "that recognition model is not installed"));
|
||||
}
|
||||
let mut st = read_settings(&dir);
|
||||
st.embedder = id;
|
||||
write_settings(&dir, &st)?;
|
||||
Ok(json!({"embedder": st.embedder}))
|
||||
}
|
||||
|
||||
fn enable(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.enable";
|
||||
let dir = models_dir(s, C)?;
|
||||
let mut st = read_settings(&dir);
|
||||
st.enabled = p.get("enabled").and_then(Value::as_bool).unwrap_or(!st.enabled);
|
||||
write_settings(&dir, &st)?;
|
||||
Ok(json!({"enabled": st.enabled}))
|
||||
}
|
||||
|
||||
pub fn specs() -> Vec<CommandSpec> {
|
||||
vec![
|
||||
cmd!(query "faces.models.list", "Face Models", [], None, "{} → {dir, enabled, embedder, runtime, models: [{id, name, role, licence{name, commercial, url, notice}, provenance, source, sizeBytes, known, installed, selected}]}", always, list),
|
||||
cmd!(query "faces.models.inspect", "Inspect Face Model File", [], None, "{path} → what a .onnx file is: {kind: known | draft | unsupported, model, assumptions, reason, alreadyInstalled}; installs nothing", always, inspect),
|
||||
cmd!(query "faces.models.install", "Install Face Model", [], None, "{path, acknowledged: true} — copy a .onnx face recognition model into the models folder. `acknowledged` must be true: the user has been shown its licence (see inspect) and accepted it", always, install),
|
||||
cmd!(query "faces.models.remove", "Remove Face Model", [], None, "{id} — delete an installed model", always, remove),
|
||||
cmd!(query "faces.models.select", "Choose Face Recognition Model", [], None, "{id: installed recogniser | null}", always, select),
|
||||
cmd!(query "faces.enable", "Face Recognition On/Off", [], None, "{enabled?: bool} (toggles when omitted)", always, enable),
|
||||
]
|
||||
}
|
||||
@@ -13,6 +13,7 @@ pub mod curves;
|
||||
mod develop;
|
||||
mod edit;
|
||||
mod export;
|
||||
mod face_models;
|
||||
pub mod filters;
|
||||
pub mod keywords;
|
||||
pub mod library;
|
||||
@@ -140,6 +141,7 @@ pub fn command_specs() -> &'static [CommandSpec] {
|
||||
v.extend(missing::specs());
|
||||
v.extend(metadata::specs());
|
||||
v.extend(filters::specs());
|
||||
v.extend(face_models::specs());
|
||||
v
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
//! Where LightCraft keeps per-user data on this machine: the config folder (app settings, face models).
|
||||
//!
|
||||
//! Hosts that have a file system (the desktop app, the CLI, the MCP server) share these defaults so a model
|
||||
//! installed from one is there for the others. Nothing here is applied automatically: a [`Session`] has no
|
||||
//! face-models folder until a host asks for one ([`Session::with_default_face_models`]), so tests stay
|
||||
//! hermetic.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crate::Session;
|
||||
|
||||
/// LightCraft's per-user config folder (`…/LightCraft`), if the platform tells us where.
|
||||
pub fn config_dir() -> Option<PathBuf> {
|
||||
if cfg!(target_os = "macos") {
|
||||
std::env::var_os("HOME").map(|h| PathBuf::from(h).join("Library/Application Support/LightCraft"))
|
||||
} else if cfg!(windows) {
|
||||
std::env::var_os("APPDATA").map(|a| PathBuf::from(a).join("LightCraft"))
|
||||
} else {
|
||||
std::env::var_os("XDG_CONFIG_HOME")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".config")))
|
||||
.map(|c| c.join("lightcraft"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Where face models are kept: `$LIGHTCRAFT_FACE_MODELS` if set, else `<config>/models`.
|
||||
pub fn default_face_models_dir() -> Option<PathBuf> {
|
||||
std::env::var_os("LIGHTCRAFT_FACE_MODELS").filter(|v| !v.is_empty()).map(PathBuf::from).or_else(|| config_dir().map(|d| d.join("models")))
|
||||
}
|
||||
|
||||
impl Session {
|
||||
/// Keep face models in the shared default folder ([`default_face_models_dir`]).
|
||||
pub fn with_default_face_models(mut self) -> Self {
|
||||
self.face_models_dir = default_face_models_dir();
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,7 @@ pub mod availability;
|
||||
mod camera_preview;
|
||||
pub mod camera_profiles;
|
||||
pub mod cmd;
|
||||
pub mod config;
|
||||
pub mod crs;
|
||||
pub mod crs_masks;
|
||||
pub mod demo;
|
||||
@@ -157,6 +158,8 @@ pub struct Session {
|
||||
/// Set by a command whose change must not rewrite the photo's XMP sidecar even with auto-write on
|
||||
/// (a catalog-only edit of data the sidecar writer does not emit); consumed when the command ends.
|
||||
pub(crate) skip_auto_write: bool,
|
||||
/// Where the host keeps face models (one folder each); `None` where there is no file system (the web).
|
||||
pub face_models_dir: Option<std::path::PathBuf>,
|
||||
/// Copied develop settings (partial JSON) for Paste.
|
||||
pub clipboard: Option<Value>,
|
||||
/// The folder on disk the [`LibrarySource::Folder`] view browses.
|
||||
@@ -263,6 +266,7 @@ impl Session {
|
||||
redo: Vec::new(),
|
||||
interaction: None,
|
||||
skip_auto_write: false,
|
||||
face_models_dir: None,
|
||||
clipboard: None,
|
||||
meta_clipboard: None,
|
||||
browse: None,
|
||||
@@ -785,6 +789,8 @@ mod tests_color;
|
||||
#[cfg(test)]
|
||||
mod tests_export;
|
||||
#[cfg(test)]
|
||||
mod tests_face_models;
|
||||
#[cfg(test)]
|
||||
mod tests_folders;
|
||||
#[cfg(test)]
|
||||
mod tests_forget_local;
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
//! The `faces.*` model commands: list, inspect, install behind the licence, select, remove, the on/off
|
||||
//! setting, and the refusals (no folder, licence not accepted, unsupported or hostile files and ids).
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::Session;
|
||||
|
||||
fn temp(name: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("lc-facemodels-{name}-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
fn session(dir: &std::path::Path) -> Session {
|
||||
let mut s = Session::new();
|
||||
s.face_models_dir = Some(dir.join("models"));
|
||||
s
|
||||
}
|
||||
|
||||
fn model_file(dir: &std::path::Path, name: &str, dim: u64) -> String {
|
||||
let p = dir.join(name);
|
||||
std::fs::write(&p, lightcraft_faces::synthetic::embedder_model(dim)).unwrap();
|
||||
p.to_string_lossy().into_owned()
|
||||
}
|
||||
|
||||
fn find<'a>(list: &'a Value, id: &str) -> &'a Value {
|
||||
list["models"].as_array().unwrap().iter().find(|m| m["id"] == id).unwrap_or(&Value::Null)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_list_knows_the_known_models_and_is_honest_about_the_runtime() {
|
||||
let d = temp("list");
|
||||
let mut s = session(&d);
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!(l["enabled"], false);
|
||||
assert_eq!(l["runtime"], false);
|
||||
assert_eq!(find(&l, "yunet-2023mar")["installed"], false, "nothing ships with LightCraft");
|
||||
assert_eq!(find(&l, "auraface-v1")["installed"], false);
|
||||
assert_eq!(find(&l, "auraface-v1")["licence"]["commercial"], "yes");
|
||||
assert_eq!(find(&l, "sface-2021dec")["licence"]["commercial"], "unknown");
|
||||
// a build with no folder lists too, and cannot install
|
||||
let mut web = Session::new();
|
||||
assert!(web.execute("faces.models.list", &json!({})).is_ok());
|
||||
let f = model_file(&d, "m.onnx", 512);
|
||||
assert!(web.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).is_err());
|
||||
assert!(web.execute("faces.enable", &json!({"enabled": true})).is_err());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inspect_describes_a_file_and_installs_nothing() {
|
||||
let d = temp("inspect");
|
||||
let mut s = session(&d);
|
||||
let f = model_file(&d, "My Model.onnx", 512);
|
||||
let r = s.execute("faces.models.inspect", &json!({"path": f})).unwrap();
|
||||
assert_eq!(r["kind"], "draft");
|
||||
assert_eq!(r["model"]["licence"]["commercial"], "unknown");
|
||||
assert_eq!(r["model"]["known"], false);
|
||||
assert!(r["assumptions"].as_array().unwrap().iter().any(|a| a.as_str().unwrap().contains("127.5")));
|
||||
assert_eq!(r["alreadyInstalled"], false);
|
||||
assert!(!d.join("models").exists(), "inspecting writes nothing");
|
||||
// not a model: a plain answer, not an error
|
||||
let junk = d.join("junk.onnx");
|
||||
std::fs::write(&junk, b"this is not a model").unwrap();
|
||||
let r = s.execute("faces.models.inspect", &json!({"path": junk.to_string_lossy()})).unwrap();
|
||||
assert_eq!(r["kind"], "unsupported");
|
||||
assert!(r["reason"].as_str().unwrap().len() > 5);
|
||||
// missing and empty files and folders are errors
|
||||
assert!(s.execute("faces.models.inspect", &json!({"path": d.join("nope.onnx").to_string_lossy()})).is_err());
|
||||
assert!(s.execute("faces.models.inspect", &json!({"path": d.to_string_lossy()})).is_err());
|
||||
let empty = d.join("empty.onnx");
|
||||
std::fs::write(&empty, b"").unwrap();
|
||||
assert!(s.execute("faces.models.inspect", &json!({"path": empty.to_string_lossy()})).is_err());
|
||||
assert!(s.execute("faces.models.inspect", &json!({})).is_err());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn install_needs_the_licence_accepted_then_select_enable_and_remove_work() {
|
||||
let d = temp("flow");
|
||||
let mut s = session(&d);
|
||||
let f = model_file(&d, "Mine.onnx", 512);
|
||||
assert!(s.execute("faces.models.install", &json!({"path": f})).is_err(), "no acknowledgement");
|
||||
assert!(s.execute("faces.models.install", &json!({"path": f, "acknowledged": false})).is_err());
|
||||
assert!(s.execute("faces.models.install", &json!({"path": f, "acknowledged": "yes"})).is_err(), "must be the boolean true");
|
||||
assert!(!d.join("models").exists(), "refused installs leave nothing behind");
|
||||
|
||||
let r = s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap();
|
||||
let id = r["installed"]["id"].as_str().unwrap().to_string();
|
||||
assert!(id.starts_with("custom-mine-"));
|
||||
let home = d.join("models").join(&id);
|
||||
assert!(home.join("model.onnx").is_file() && home.join("face-model.json").is_file() && home.join("installed.json").is_file());
|
||||
assert!(!home.join("model.onnx.part").exists());
|
||||
assert_eq!(std::fs::read(home.join("model.onnx")).unwrap(), std::fs::read(&f).unwrap());
|
||||
let rec: Value = serde_json::from_slice(&std::fs::read(home.join("installed.json")).unwrap()).unwrap();
|
||||
assert_eq!(rec["commercial"], "unknown");
|
||||
assert_eq!(rec["fileName"], "Mine.onnx");
|
||||
|
||||
// listed as installed and custom; installing again is harmless
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!(find(&l, &id)["installed"], true);
|
||||
assert_eq!(find(&l, &id)["known"], false);
|
||||
s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap();
|
||||
assert_eq!(s.execute("faces.models.inspect", &json!({"path": f})).unwrap()["alreadyInstalled"], true);
|
||||
|
||||
// choose it; the choice and the on/off setting survive a new session on the same folder
|
||||
assert!(s.execute("faces.models.select", &json!({"id": "auraface-v1"})).is_err(), "not installed");
|
||||
assert!(s.execute("faces.models.select", &json!({"id": "yunet-2023mar"})).is_err(), "a detector is not a recogniser");
|
||||
s.execute("faces.models.select", &json!({"id": id})).unwrap();
|
||||
assert_eq!(s.execute("faces.enable", &json!({"enabled": true})).unwrap()["enabled"], true);
|
||||
let mut again = session(&d);
|
||||
let l = again.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!((l["enabled"].clone(), l["embedder"].clone()), (json!(true), json!(id)));
|
||||
assert_eq!(again.execute("faces.enable", &json!({})).unwrap()["enabled"], false, "toggles when omitted");
|
||||
|
||||
// removing it clears the choice
|
||||
assert_eq!(s.execute("faces.models.remove", &json!({"id": id})).unwrap()["removed"], id);
|
||||
assert!(!home.exists());
|
||||
assert_eq!(s.execute("faces.models.list", &json!({})).unwrap()["embedder"], Value::Null);
|
||||
assert!(s.execute("faces.models.remove", &json!({"id": id})).is_err(), "already gone");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hostile_ids_unsupported_models_and_odd_folders_are_handled() {
|
||||
let d = temp("hostile");
|
||||
let mut s = session(&d);
|
||||
let models = d.join("models");
|
||||
std::fs::create_dir_all(&models).unwrap();
|
||||
// something that is not a model sits in the folder: it is never listed and never deleted
|
||||
let precious = models.join("precious");
|
||||
std::fs::create_dir_all(&precious).unwrap();
|
||||
std::fs::write(precious.join("keep.txt"), b"mine").unwrap();
|
||||
for id in ["precious", "..", ".", "", "../x", "a/b", r"a\b", ".hidden", "yunet-2023mar", "UPPER"] {
|
||||
assert!(s.execute("faces.models.remove", &json!({"id": id})).is_err(), "{id:?}");
|
||||
}
|
||||
assert!(precious.join("keep.txt").is_file());
|
||||
assert!(d.join("models").is_dir());
|
||||
// a folder whose manifest lies about its id is ignored
|
||||
let liar = models.join("liar");
|
||||
std::fs::create_dir_all(&liar).unwrap();
|
||||
std::fs::write(liar.join("model.onnx"), lightcraft_faces::synthetic::embedder_model(64)).unwrap();
|
||||
let mut m = lightcraft_faces::known::auraface();
|
||||
m.id = "someone-else".into();
|
||||
std::fs::write(liar.join("face-model.json"), serde_json::to_vec(&m).unwrap()).unwrap();
|
||||
std::fs::write(models.join("settings.json"), b"{ not json").unwrap();
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert!(find(&l, "someone-else").is_null() && find(&l, "liar").is_null());
|
||||
assert_eq!(l["enabled"], false, "a corrupt settings file is ignored");
|
||||
// not-a-model and wrong-shaped files cannot be installed
|
||||
let junk = d.join("junk.onnx");
|
||||
std::fs::write(&junk, b"nonsense").unwrap();
|
||||
let e = s.execute("faces.models.install", &json!({"path": junk.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("cannot be used yet"), "{e}");
|
||||
let tokens = d.join("tokens.onnx");
|
||||
{
|
||||
use lightcraft_faces::synthetic::{len_field, value_info_bytes};
|
||||
let mut graph = Vec::new();
|
||||
len_field(11, &value_info_bytes("pixels", 1, &[Err("N"), Ok(3), Ok(224), Ok(224)]), &mut graph);
|
||||
len_field(12, &value_info_bytes("tokens", 1, &[Err("N"), Ok(257), Ok(384)]), &mut graph);
|
||||
let mut model = Vec::new();
|
||||
len_field(7, &graph, &mut model);
|
||||
std::fs::write(&tokens, model).unwrap();
|
||||
}
|
||||
assert!(s.execute("faces.models.install", &json!({"path": tokens.to_string_lossy(), "acknowledged": true})).is_err());
|
||||
assert!(s.execute("faces.models.select", &json!({"id": 5})).is_err());
|
||||
assert!(s.execute("faces.models.select", &json!({"id": ""})).is_err());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
[package]
|
||||
name = "lightcraft-faces"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
description = "LightCraft faces (L1): face-model manifests, validation, the known-model table and an ONNX shape probe"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
proptest = { workspace = true }
|
||||
@@ -0,0 +1,46 @@
|
||||
//! SHA-256 of a model file, streamed (a model can be hundreds of megabytes).
|
||||
|
||||
use std::io::{self, Read};
|
||||
use std::path::Path;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// Lowercase hex SHA-256 of everything `r` yields.
|
||||
pub fn sha256_reader<R: Read>(mut r: R) -> io::Result<String> {
|
||||
let mut h = Sha256::new();
|
||||
let mut buf = vec![0u8; 1 << 20];
|
||||
loop {
|
||||
let n = r.read(&mut buf)?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
h.update(buf.get(..n).unwrap_or(&[]));
|
||||
}
|
||||
Ok(hex(&h.finalize()))
|
||||
}
|
||||
|
||||
/// Lowercase hex SHA-256 of a file.
|
||||
pub fn sha256_file(path: &Path) -> io::Result<String> {
|
||||
sha256_reader(io::BufReader::new(std::fs::File::open(path)?))
|
||||
}
|
||||
|
||||
/// Lowercase hex SHA-256 of `bytes`.
|
||||
pub fn sha256_hex(bytes: &[u8]) -> String {
|
||||
hex(&Sha256::digest(bytes))
|
||||
}
|
||||
|
||||
fn hex(bytes: &[u8]) -> String {
|
||||
bytes.iter().map(|b| format!("{b:02x}")).collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn known_digests() {
|
||||
assert_eq!(sha256_hex(b""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
|
||||
assert_eq!(sha256_hex(b"abc"), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");
|
||||
assert_eq!(sha256_reader(&b"abc"[..]).unwrap(), sha256_hex(b"abc"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
//! The models LightCraft recognises by their SHA-256, with what is honestly known about each.
|
||||
//!
|
||||
//! The licence and provenance texts here are what the user reads before enabling a model, so they say what
|
||||
//! is *not* known too. Weights are never part of LightCraft: every model is an opt-in the user installs.
|
||||
|
||||
use crate::manifest::{Colour, Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Resize, Role, Thresholds};
|
||||
|
||||
/// Detector output decoders built into LightCraft.
|
||||
pub const DECODERS: &[&str] = &["yunet-v2"];
|
||||
|
||||
pub const YUNET_SHA256: &str = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4";
|
||||
pub const SFACE_SHA256: &str = "0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79";
|
||||
pub const AURAFACE_SHA256: &str = "a7933ea5330113b01c9b60351d8f4c33003f145d8470ac5f0e52ee2effe25c60";
|
||||
|
||||
fn licence(name: &str, commercial: Commercial, url: &str, notice: &str) -> Licence {
|
||||
Licence { name: name.into(), commercial, url: Some(url.into()), notice: notice.into() }
|
||||
}
|
||||
|
||||
/// YuNet 2023mar (OpenCV Zoo): a small face detector.
|
||||
pub fn yunet() -> ModelManifest {
|
||||
ModelManifest {
|
||||
id: "yunet-2023mar".into(),
|
||||
name: "YuNet (face detector)".into(),
|
||||
version: "2023mar".into(),
|
||||
role: Role::Detector,
|
||||
licence: licence(
|
||||
"MIT",
|
||||
Commercial::Yes,
|
||||
"https://github.com/opencv/opencv_zoo/blob/main/models/face_detection_yunet/LICENSE",
|
||||
"MIT licence, copyright Shiqi Yu.",
|
||||
),
|
||||
source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet".into()),
|
||||
sha256: Some(YUNET_SHA256.into()),
|
||||
size_bytes: Some(232_589),
|
||||
provenance:
|
||||
"Trained on WIDER FACE, whose terms are not settled by the weights' MIT licence. A detector: it finds faces, it does not identify anyone."
|
||||
.into(),
|
||||
input: InputSpec { width: 640, height: 640, colour: Colour::Bgr, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Letterbox },
|
||||
output: OutputSpec::Detector { decoder: "yunet-v2".into() },
|
||||
thresholds: Thresholds { match_cosine: None, score: Some(0.6), nms_iou: Some(0.3) },
|
||||
}
|
||||
}
|
||||
|
||||
/// SFace 2021dec (OpenCV Zoo).
|
||||
pub fn sface() -> ModelManifest {
|
||||
ModelManifest {
|
||||
id: "sface-2021dec".into(),
|
||||
name: "SFace (face recogniser)".into(),
|
||||
version: "2021dec".into(),
|
||||
role: Role::Embedder,
|
||||
licence: licence(
|
||||
"Apache-2.0 (as labelled by OpenCV Zoo)",
|
||||
Commercial::Unknown,
|
||||
"https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface",
|
||||
"Labelled Apache-2.0, but what it was trained on is not documented, and two questions about commercial use (opencv_zoo issues 313 and 318) are unanswered. Fine to try for yourself; do not redistribute.",
|
||||
),
|
||||
source: Some("https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface".into()),
|
||||
sha256: Some(SFACE_SHA256.into()),
|
||||
size_bytes: Some(38_696_353),
|
||||
provenance: "Undocumented. The original SFace repository mentions CASIA-WebFace, VGGFace2 and MS1MV2.".into(),
|
||||
input: InputSpec { width: 112, height: 112, colour: Colour::Rgb, mean: [0.0; 3], std: [1.0; 3], resize: Resize::Stretch },
|
||||
output: OutputSpec::Embedding { dim: 128 },
|
||||
thresholds: Thresholds::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// AuraFace v1 `glintr100` (fal.ai).
|
||||
pub fn auraface() -> ModelManifest {
|
||||
ModelManifest {
|
||||
id: "auraface-v1".into(),
|
||||
name: "AuraFace v1 (face recogniser)".into(),
|
||||
version: "1".into(),
|
||||
role: Role::Embedder,
|
||||
licence: licence(
|
||||
"Apache-2.0",
|
||||
Commercial::Yes,
|
||||
"https://huggingface.co/fal/AuraFace-v1",
|
||||
"Apache-2.0. Its training data is described only as a commercial dataset, with no dataset named and no consent statement, and it covers some ethnicities less well. About 261 MB, and slow on a CPU.",
|
||||
),
|
||||
source: Some("https://huggingface.co/fal/AuraFace-v1".into()),
|
||||
sha256: Some(AURAFACE_SHA256.into()),
|
||||
size_bytes: Some(260_694_151),
|
||||
provenance: "Undisclosed: \"a commercial dataset comprising face images from various sources\".".into(),
|
||||
input: InputSpec::default(),
|
||||
output: OutputSpec::Embedding { dim: 512 },
|
||||
thresholds: Thresholds::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every model we know by hash.
|
||||
pub fn all() -> Vec<ModelManifest> {
|
||||
vec![yunet(), sface(), auraface()]
|
||||
}
|
||||
|
||||
/// The known model whose file has this SHA-256 (lowercase hex).
|
||||
pub fn lookup(sha256: &str) -> Option<ModelManifest> {
|
||||
all().into_iter().find(|m| m.sha256.as_deref() == Some(sha256))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::manifest::validate;
|
||||
|
||||
#[test]
|
||||
fn known_models_are_valid_and_unique() {
|
||||
let all = all();
|
||||
for m in &all {
|
||||
assert_eq!(validate(m), Ok(()), "{}", m.id);
|
||||
}
|
||||
let mut ids: Vec<_> = all.iter().map(|m| &m.id).collect();
|
||||
ids.sort();
|
||||
ids.dedup();
|
||||
assert_eq!(ids.len(), all.len(), "unique ids");
|
||||
let mut hashes: Vec<_> = all.iter().filter_map(|m| m.sha256.as_deref()).collect();
|
||||
hashes.sort();
|
||||
hashes.dedup();
|
||||
assert_eq!(hashes.len(), all.len(), "unique hashes");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lookup_finds_by_hash() {
|
||||
assert_eq!(lookup(AURAFACE_SHA256).map(|m| m.id), Some("auraface-v1".to_string()));
|
||||
assert_eq!(lookup(YUNET_SHA256).map(|m| m.role), Some(Role::Detector));
|
||||
assert!(lookup(&"0".repeat(64)).is_none());
|
||||
assert!(lookup("").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unresolved_models_say_so() {
|
||||
assert_eq!(sface().licence.commercial, Commercial::Unknown);
|
||||
assert!(sface().provenance.to_lowercase().contains("undocumented"));
|
||||
assert!(auraface().provenance.to_lowercase().contains("undisclosed"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! Face models for LightCraft: what a model is ([`ModelManifest`]), whether to trust a description of one
|
||||
//! ([`manifest::validate`]), the models we know by hash ([`known`]), and a probe that reads an unknown
|
||||
//! `.onnx` file's input and output shapes without running it ([`onnx`]), so installing a model of your
|
||||
//! own can be one drop and one licence prompt.
|
||||
//!
|
||||
//! Everything here treats model files and manifests as hostile input: sizes are capped, numbers must be
|
||||
//! finite and sane, and a malformed file is an error, never a panic. Model execution is not part of
|
||||
//! this crate yet.
|
||||
#![forbid(unsafe_code)]
|
||||
#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)]
|
||||
|
||||
pub mod hash;
|
||||
pub mod known;
|
||||
pub mod manifest;
|
||||
pub mod onnx;
|
||||
pub mod suggest;
|
||||
pub mod synthetic;
|
||||
|
||||
pub use manifest::{Colour, Commercial, InputSpec, Licence, ManifestError, ModelManifest, OutputSpec, Resize, Role, Thresholds};
|
||||
@@ -0,0 +1,344 @@
|
||||
//! A face model's manifest: what the model is, what it needs and who may use it, as plain data.
|
||||
//!
|
||||
//! A manifest comes from the known-model table, from a suggestion made for a file the user dropped in, or
|
||||
//! from a `face-model.json` next to a model. All of those are untrusted: [`validate`] must pass before a
|
||||
//! manifest is used, and [`parse`] applies it after reading JSON of bounded size.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Largest `face-model.json` we read.
|
||||
pub const MAX_MANIFEST_BYTES: usize = 64 * 1024;
|
||||
/// Largest model file we accept (bytes).
|
||||
pub const MAX_MODEL_BYTES: u64 = 4 * 1024 * 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum Role {
|
||||
/// Finds faces (boxes, and landmarks when the model has them).
|
||||
Detector,
|
||||
/// Turns an aligned face into a vector whose distance says whether two faces are the same person.
|
||||
Embedder,
|
||||
}
|
||||
|
||||
/// Whether the weights may be used commercially.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum Commercial {
|
||||
Yes,
|
||||
/// Research or personal use only. Never bundled or redistributed by LightCraft.
|
||||
No,
|
||||
/// The licence or the training data is unclear: the user is told so before installing.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default, rename_all = "camelCase")]
|
||||
pub struct Licence {
|
||||
/// SPDX id or a short name ("MIT", "Apache-2.0", "InsightFace non-commercial research").
|
||||
pub name: String,
|
||||
pub commercial: Commercial,
|
||||
pub url: Option<String>,
|
||||
/// Plain-language terms shown before the model is enabled.
|
||||
pub notice: String,
|
||||
}
|
||||
|
||||
impl Default for Licence {
|
||||
fn default() -> Self {
|
||||
Self { name: String::new(), commercial: Commercial::Unknown, url: None, notice: String::new() }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum Colour {
|
||||
Rgb,
|
||||
Bgr,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum Resize {
|
||||
/// Stretch the crop to the input size (aligned face crops).
|
||||
Stretch,
|
||||
/// Keep the aspect ratio and pad (whole photos for a detector).
|
||||
Letterbox,
|
||||
}
|
||||
|
||||
/// How an image becomes the model's input: NCHW float, `(pixel − mean) / std` per channel, pixel in 0..255.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default, rename_all = "camelCase")]
|
||||
pub struct InputSpec {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub colour: Colour,
|
||||
pub mean: [f32; 3],
|
||||
pub std: [f32; 3],
|
||||
pub resize: Resize,
|
||||
}
|
||||
|
||||
impl Default for InputSpec {
|
||||
/// ArcFace-style: 112 × 112, RGB, `(x − 127.5) / 127.5`.
|
||||
fn default() -> Self {
|
||||
Self { width: 112, height: 112, colour: Colour::Rgb, mean: [127.5; 3], std: [127.5; 3], resize: Resize::Stretch }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "camelCase")]
|
||||
pub enum OutputSpec {
|
||||
/// One vector of `dim` numbers per face.
|
||||
Embedding { dim: u32 },
|
||||
/// A detector whose outputs a built-in decoder understands (`"yunet-v2"`).
|
||||
Detector { decoder: String },
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default, rename_all = "camelCase")]
|
||||
pub struct Thresholds {
|
||||
/// Embedders: cosine similarity at or above which two faces are suggested as a match.
|
||||
pub match_cosine: Option<f32>,
|
||||
/// Detectors: lowest detection score kept.
|
||||
pub score: Option<f32>,
|
||||
/// Detectors: overlap above which two boxes are one face.
|
||||
pub nms_iou: Option<f32>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ModelManifest {
|
||||
/// Lowercase letters, digits, `.` `_` `-`.
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
pub role: Role,
|
||||
#[serde(default)]
|
||||
pub licence: Licence,
|
||||
/// Where the weights come from (an `https://` page or file), for the user's information.
|
||||
#[serde(default)]
|
||||
pub source: Option<String>,
|
||||
/// Lowercase hex SHA-256 of the `.onnx` file, when known.
|
||||
#[serde(default)]
|
||||
pub sha256: Option<String>,
|
||||
#[serde(default)]
|
||||
pub size_bytes: Option<u64>,
|
||||
/// What the model was trained on, in plain words; says "undisclosed" when it is.
|
||||
#[serde(default)]
|
||||
pub provenance: String,
|
||||
#[serde(default)]
|
||||
pub input: InputSpec,
|
||||
pub output: OutputSpec,
|
||||
#[serde(default)]
|
||||
pub thresholds: Thresholds,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, thiserror::Error)]
|
||||
pub enum ManifestError {
|
||||
#[error("`{0}` is not valid: {1}")]
|
||||
Field(&'static str, String),
|
||||
#[error("not a face-model manifest: {0}")]
|
||||
Parse(String),
|
||||
#[error("a manifest may be at most {MAX_MANIFEST_BYTES} bytes")]
|
||||
TooLarge,
|
||||
}
|
||||
|
||||
fn bad<T>(field: &'static str, why: impl Into<String>) -> Result<T, ManifestError> {
|
||||
Err(ManifestError::Field(field, why.into()))
|
||||
}
|
||||
|
||||
/// A model id is a folder name: 1 to 64 characters, starting with a lowercase letter or digit, then
|
||||
/// lowercase letters, digits, `.`, `_` or `-`. So it is never `.`, `..`, hidden, or a path.
|
||||
pub fn valid_id(id: &str) -> bool {
|
||||
let mut bytes = id.bytes();
|
||||
let first_ok = bytes.next().is_some_and(|b| b.is_ascii_lowercase() || b.is_ascii_digit());
|
||||
first_ok && id.len() <= 64 && bytes.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || matches!(b, b'.' | b'_' | b'-'))
|
||||
}
|
||||
|
||||
fn text(field: &'static str, s: &str, min: usize, max: usize) -> Result<(), ManifestError> {
|
||||
let n = s.chars().count();
|
||||
if n < min || n > max {
|
||||
return bad(field, format!("must be {min} to {max} characters"));
|
||||
}
|
||||
if s.chars().any(|c| c.is_control() && c != '\n') {
|
||||
return bad(field, "must not contain control characters");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn url(field: &'static str, u: &Option<String>) -> Result<(), ManifestError> {
|
||||
match u {
|
||||
None => Ok(()),
|
||||
Some(u) if u.len() <= 500 && u.starts_with("https://") && !u.chars().any(|c| c.is_control() || c.is_whitespace()) => Ok(()),
|
||||
Some(_) => bad(field, "must be an https:// address of at most 500 characters"),
|
||||
}
|
||||
}
|
||||
|
||||
fn ratio(field: &'static str, v: Option<f32>, lo: f32, hi: f32) -> Result<(), ManifestError> {
|
||||
match v {
|
||||
Some(v) if !(v.is_finite() && (lo..=hi).contains(&v)) => bad(field, format!("must be a number from {lo} to {hi}")),
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check a manifest from any source. Everything the rest of LightCraft relies on (sizes it will allocate,
|
||||
/// numbers it will divide by, text it will show) is bounded here.
|
||||
pub fn validate(m: &ModelManifest) -> Result<(), ManifestError> {
|
||||
if !valid_id(&m.id) {
|
||||
return bad("id", "must be 1 to 64 characters: lowercase letters and digits, then also '.', '_' or '-'");
|
||||
}
|
||||
text("name", &m.name, 1, 120)?;
|
||||
text("version", &m.version, 1, 40)?;
|
||||
text("licence.name", &m.licence.name, 0, 80)?;
|
||||
text("licence.notice", &m.licence.notice, 0, 800)?;
|
||||
text("provenance", &m.provenance, 0, 800)?;
|
||||
url("licence.url", &m.licence.url)?;
|
||||
url("source", &m.source)?;
|
||||
if let Some(h) = &m.sha256
|
||||
&& !(h.len() == 64 && h.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)))
|
||||
{
|
||||
return bad("sha256", "must be 64 lowercase hexadecimal digits");
|
||||
}
|
||||
if m.size_bytes.is_some_and(|s| s == 0 || s > MAX_MODEL_BYTES) {
|
||||
return bad("sizeBytes", format!("must be 1 to {MAX_MODEL_BYTES}"));
|
||||
}
|
||||
let i = &m.input;
|
||||
if !(16..=4096).contains(&i.width) || !(16..=4096).contains(&i.height) {
|
||||
return bad("input.width/height", "must be 16 to 4096");
|
||||
}
|
||||
if i.mean.iter().any(|v| !v.is_finite() || v.abs() > 1000.0) {
|
||||
return bad("input.mean", "must be finite numbers within ±1000");
|
||||
}
|
||||
if i.std.iter().any(|v| !v.is_finite() || v.abs() < 1e-3 || v.abs() > 1000.0) {
|
||||
return bad("input.std", "must be finite, not zero, within ±1000");
|
||||
}
|
||||
match &m.output {
|
||||
OutputSpec::Embedding { dim } => {
|
||||
if m.role != Role::Embedder {
|
||||
return bad("output", "an embedding output needs the embedder role");
|
||||
}
|
||||
if !(1..=4096).contains(dim) {
|
||||
return bad("output.dim", "must be 1 to 4096");
|
||||
}
|
||||
}
|
||||
OutputSpec::Detector { decoder } => {
|
||||
if m.role != Role::Detector {
|
||||
return bad("output", "a detector output needs the detector role");
|
||||
}
|
||||
if !crate::known::DECODERS.contains(&decoder.as_str()) {
|
||||
return bad("output.decoder", format!("unknown decoder `{}`", decoder.chars().take(40).collect::<String>()));
|
||||
}
|
||||
}
|
||||
}
|
||||
ratio("thresholds.matchCosine", m.thresholds.match_cosine, -1.0, 1.0)?;
|
||||
ratio("thresholds.score", m.thresholds.score, 0.0, 1.0)?;
|
||||
ratio("thresholds.nmsIou", m.thresholds.nms_iou, 0.0, 1.0)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read a `face-model.json` (at most [`MAX_MANIFEST_BYTES`]) and validate it.
|
||||
pub fn parse(bytes: &[u8]) -> Result<ModelManifest, ManifestError> {
|
||||
if bytes.len() > MAX_MANIFEST_BYTES {
|
||||
return Err(ManifestError::TooLarge);
|
||||
}
|
||||
let m: ModelManifest = serde_json::from_slice(bytes).map_err(|e| ManifestError::Parse(e.to_string()))?;
|
||||
validate(&m)?;
|
||||
Ok(m)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn good() -> ModelManifest {
|
||||
ModelManifest {
|
||||
id: "example-r50".into(),
|
||||
name: "Example R50".into(),
|
||||
version: "1".into(),
|
||||
role: Role::Embedder,
|
||||
licence: Licence {
|
||||
name: "MIT".into(),
|
||||
commercial: Commercial::Yes,
|
||||
url: Some("https://example.org/licence".into()),
|
||||
notice: String::new(),
|
||||
},
|
||||
source: Some("https://example.org/model.onnx".into()),
|
||||
sha256: Some("a".repeat(64)),
|
||||
size_bytes: Some(1234),
|
||||
provenance: "synthetic".into(),
|
||||
input: InputSpec::default(),
|
||||
output: OutputSpec::Embedding { dim: 512 },
|
||||
thresholds: Thresholds { match_cosine: Some(0.4), ..Default::default() },
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_good_manifest_passes_and_round_trips() {
|
||||
let m = good();
|
||||
assert_eq!(validate(&m), Ok(()));
|
||||
let json = serde_json::to_vec(&m).unwrap();
|
||||
assert_eq!(parse(&json).unwrap(), m);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_field_is_bounded() {
|
||||
let cases: Vec<(&str, Box<dyn Fn(&mut ModelManifest)>)> = vec![
|
||||
("empty id", Box::new(|m| m.id = String::new())),
|
||||
("uppercase id", Box::new(|m| m.id = "Bad".into())),
|
||||
("path in id", Box::new(|m| m.id = "../etc".into())),
|
||||
("dot id", Box::new(|m| m.id = ".".into())),
|
||||
("dotdot id", Box::new(|m| m.id = "..".into())),
|
||||
("hidden id", Box::new(|m| m.id = ".hidden".into())),
|
||||
("slash id", Box::new(|m| m.id = "a/b".into())),
|
||||
("backslash id", Box::new(|m| m.id = r"a\b".into())),
|
||||
("long id", Box::new(|m| m.id = "a".repeat(65))),
|
||||
("empty name", Box::new(|m| m.name = String::new())),
|
||||
("control in name", Box::new(|m| m.name = "a\u{7}b".into())),
|
||||
("huge notice", Box::new(|m| m.licence.notice = "x".repeat(801))),
|
||||
("http source", Box::new(|m| m.source = Some("http://example.org/m.onnx".into()))),
|
||||
("file source", Box::new(|m| m.source = Some("file:///etc/passwd".into()))),
|
||||
("space in url", Box::new(|m| m.source = Some("https://a b".into()))),
|
||||
("short sha", Box::new(|m| m.sha256 = Some("abc".into()))),
|
||||
("uppercase sha", Box::new(|m| m.sha256 = Some("A".repeat(64)))),
|
||||
("zero size", Box::new(|m| m.size_bytes = Some(0))),
|
||||
("huge size", Box::new(|m| m.size_bytes = Some(u64::MAX))),
|
||||
("tiny input", Box::new(|m| m.input.width = 4)),
|
||||
("huge input", Box::new(|m| m.input.height = 1 << 20)),
|
||||
("nan mean", Box::new(|m| m.input.mean[1] = f32::NAN)),
|
||||
("inf mean", Box::new(|m| m.input.mean[0] = f32::INFINITY)),
|
||||
("zero std", Box::new(|m| m.input.std[2] = 0.0)),
|
||||
("nan std", Box::new(|m| m.input.std[0] = f32::NAN)),
|
||||
("zero dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 0 })),
|
||||
("huge dim", Box::new(|m| m.output = OutputSpec::Embedding { dim: 10_000_000 })),
|
||||
("role mismatch", Box::new(|m| m.role = Role::Detector)),
|
||||
(
|
||||
"unknown decoder",
|
||||
Box::new(|m| {
|
||||
m.role = Role::Detector;
|
||||
m.output = OutputSpec::Detector { decoder: "made-up".into() };
|
||||
}),
|
||||
),
|
||||
("nan threshold", Box::new(|m| m.thresholds.match_cosine = Some(f32::NAN))),
|
||||
("threshold range", Box::new(|m| m.thresholds.score = Some(1.5))),
|
||||
];
|
||||
for (what, f) in cases {
|
||||
let mut m = good();
|
||||
f(&mut m);
|
||||
assert!(validate(&m).is_err(), "{what} must be rejected");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hostile_json_is_an_error() {
|
||||
assert!(matches!(parse(&vec![b' '; MAX_MANIFEST_BYTES + 1]), Err(ManifestError::TooLarge)));
|
||||
for bytes in [
|
||||
&b""[..],
|
||||
b"null",
|
||||
b"[]",
|
||||
b"{",
|
||||
b"{\"id\": 1}",
|
||||
b"\xff\xfe\x00",
|
||||
b"{\"id\":\"x\",\"name\":\"n\",\"version\":\"1\",\"role\":\"embedder\"}",
|
||||
] {
|
||||
assert!(parse(bytes).is_err(), "{:?}", String::from_utf8_lossy(bytes));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,465 @@
|
||||
//! A streaming probe of an ONNX file: the graph's inputs and outputs (name, element type, shape), without
|
||||
//! running or even fully loading the model.
|
||||
//!
|
||||
//! ONNX is protobuf. We read only the few fields we need (the producer, the opset, and the graph's input,
|
||||
//! output and initializer names) and seek past everything else, so a 260 MB model costs a few kilobytes and
|
||||
//! a few milliseconds. The file is hostile input: every length is checked against its parent message and
|
||||
//! the file size, the number of fields read is capped, strings are bounded, and nothing indexes or unwraps.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::io::{self, Read, Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
|
||||
const MAX_FIELDS: u32 = 5_000_000;
|
||||
const MAX_STRING: u64 = 4096;
|
||||
const MAX_DIMS: usize = 16;
|
||||
const MAX_TENSORS: usize = 64;
|
||||
/// Older exporters list every weight as a graph input too; that many are read before the weights are filtered out.
|
||||
const MAX_LISTED_INPUTS: usize = 200_000;
|
||||
const MAX_INITIALIZERS: usize = 1_000_000;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Dim {
|
||||
Fixed(u64),
|
||||
/// A named or unknown size (typically the batch).
|
||||
Dynamic(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct TensorInfo {
|
||||
pub name: String,
|
||||
/// ONNX element type (1 = float32).
|
||||
pub elem_type: u32,
|
||||
pub shape: Vec<Dim>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct OnnxInfo {
|
||||
pub producer: String,
|
||||
/// The default-domain opset version.
|
||||
pub opset: Option<u64>,
|
||||
/// Real inputs only (weights that older exporters also list as inputs are left out).
|
||||
pub inputs: Vec<TensorInfo>,
|
||||
pub outputs: Vec<TensorInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ProbeError {
|
||||
#[error("not an ONNX model: {0}")]
|
||||
Malformed(&'static str),
|
||||
#[error("could not read the model: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
}
|
||||
|
||||
type Result<T> = std::result::Result<T, ProbeError>;
|
||||
|
||||
enum Wire {
|
||||
Varint(u64),
|
||||
/// A length-delimited field: its bytes are `start..end` and the reader sits at `start`.
|
||||
Len {
|
||||
start: u64,
|
||||
end: u64,
|
||||
},
|
||||
/// Fixed-width data already skipped.
|
||||
Skipped,
|
||||
}
|
||||
|
||||
struct Walker<'a, R> {
|
||||
r: &'a mut R,
|
||||
total: u64,
|
||||
fields: u32,
|
||||
}
|
||||
|
||||
impl<R: Read + Seek> Walker<'_, R> {
|
||||
fn pos(&mut self) -> Result<u64> {
|
||||
Ok(self.r.stream_position()?)
|
||||
}
|
||||
|
||||
fn seek_to(&mut self, at: u64) -> Result<()> {
|
||||
self.r.seek(SeekFrom::Start(at))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn byte(&mut self) -> Result<u8> {
|
||||
let mut b = [0u8; 1];
|
||||
self.r
|
||||
.read_exact(&mut b)
|
||||
.map_err(|e| if e.kind() == io::ErrorKind::UnexpectedEof { ProbeError::Malformed("file ends inside a field") } else { e.into() })?;
|
||||
let [x] = b;
|
||||
Ok(x)
|
||||
}
|
||||
|
||||
fn varint(&mut self) -> Result<u64> {
|
||||
let mut v = 0u64;
|
||||
for shift in (0..70u32).step_by(7) {
|
||||
let b = self.byte()?;
|
||||
if shift < 64 {
|
||||
v |= u64::from(b & 0x7f) << shift;
|
||||
}
|
||||
if b & 0x80 == 0 {
|
||||
return Ok(v);
|
||||
}
|
||||
}
|
||||
Err(ProbeError::Malformed("number too long"))
|
||||
}
|
||||
|
||||
/// The next field of the message that ends at `end`, or `None` at its end.
|
||||
fn next(&mut self, end: u64) -> Result<Option<(u32, Wire)>> {
|
||||
let here = self.pos()?;
|
||||
if here >= end {
|
||||
return Ok(None);
|
||||
}
|
||||
self.fields += 1;
|
||||
if self.fields > MAX_FIELDS {
|
||||
return Err(ProbeError::Malformed("too many fields"));
|
||||
}
|
||||
let tag = self.varint()?;
|
||||
let number = u32::try_from(tag >> 3).map_err(|_| ProbeError::Malformed("field number too large"))?;
|
||||
let wire = match tag & 7 {
|
||||
0 => Wire::Varint(self.varint()?),
|
||||
1 | 5 => {
|
||||
let n = if tag & 7 == 1 { 8 } else { 4 };
|
||||
let at = self.pos()?.checked_add(n).ok_or(ProbeError::Malformed("offset overflow"))?;
|
||||
if at > end {
|
||||
return Err(ProbeError::Malformed("field runs past its message"));
|
||||
}
|
||||
self.seek_to(at)?;
|
||||
Wire::Skipped
|
||||
}
|
||||
2 => {
|
||||
let len = self.varint()?;
|
||||
let start = self.pos()?;
|
||||
let stop = start.checked_add(len).ok_or(ProbeError::Malformed("offset overflow"))?;
|
||||
if stop > end || stop > self.total {
|
||||
return Err(ProbeError::Malformed("field runs past its message"));
|
||||
}
|
||||
Wire::Len { start, end: stop }
|
||||
}
|
||||
_ => return Err(ProbeError::Malformed("unsupported field type")),
|
||||
};
|
||||
Ok(Some((number, wire)))
|
||||
}
|
||||
|
||||
/// A bounded UTF-8 string occupying `start..end` (empty when it is longer than we care about).
|
||||
fn string(&mut self, start: u64, end: u64) -> Result<String> {
|
||||
let len = end.saturating_sub(start);
|
||||
if len > MAX_STRING {
|
||||
return Ok(String::new());
|
||||
}
|
||||
let mut buf = Vec::new();
|
||||
self.r.by_ref().take(len).read_to_end(&mut buf)?;
|
||||
Ok(String::from_utf8_lossy(&buf).into_owned())
|
||||
}
|
||||
}
|
||||
|
||||
fn dim<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Option<Dim>> {
|
||||
let mut out = None;
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(1, Wire::Varint(v)) => out = Some(Dim::Fixed(v)),
|
||||
(2, Wire::Len { start, end: e }) => {
|
||||
out = Some(Dim::Dynamic(w.string(start, e)?));
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn shape<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Vec<Dim>> {
|
||||
let mut dims = Vec::new();
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(1, Wire::Len { end: e, .. }) => {
|
||||
if dims.len() >= MAX_DIMS {
|
||||
return Err(ProbeError::Malformed("too many dimensions"));
|
||||
}
|
||||
dims.push(dim(w, e)?.unwrap_or(Dim::Dynamic(String::new())));
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(dims)
|
||||
}
|
||||
|
||||
fn tensor_type<R: Read + Seek>(w: &mut Walker<R>, end: u64, info: &mut TensorInfo) -> Result<()> {
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(1, Wire::Varint(v)) => info.elem_type = u32::try_from(v).unwrap_or(0),
|
||||
(2, Wire::Len { end: e, .. }) => {
|
||||
info.shape = shape(w, e)?;
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn type_proto<R: Read + Seek>(w: &mut Walker<R>, end: u64, info: &mut TensorInfo) -> Result<()> {
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(1, Wire::Len { end: e, .. }) => {
|
||||
tensor_type(w, e, info)?;
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn value_info<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<TensorInfo> {
|
||||
let mut info = TensorInfo { name: String::new(), elem_type: 0, shape: Vec::new() };
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(1, Wire::Len { start, end: e }) => {
|
||||
info.name = w.string(start, e)?;
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(2, Wire::Len { end: e, .. }) => {
|
||||
type_proto(w, e, &mut info)?;
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
/// The `name` (field 8) of a weight tensor; its bulk data is skipped, never read.
|
||||
fn tensor_name<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Option<String>> {
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(8, Wire::Len { start, end: e }) => return Ok(Some(w.string(start, e)?)),
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn graph<R: Read + Seek>(w: &mut Walker<R>, end: u64, info: &mut OnnxInfo, weights: &mut HashSet<String>) -> Result<()> {
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(5, Wire::Len { end: e, .. }) => {
|
||||
if weights.len() < MAX_INITIALIZERS
|
||||
&& let Some(name) = tensor_name(w, e)?
|
||||
{
|
||||
weights.insert(name);
|
||||
}
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(11, Wire::Len { end: e, .. }) => {
|
||||
if info.inputs.len() >= MAX_LISTED_INPUTS {
|
||||
return Err(ProbeError::Malformed("too many inputs"));
|
||||
}
|
||||
info.inputs.push(value_info(w, e)?);
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(12, Wire::Len { end: e, .. }) => {
|
||||
if info.outputs.len() >= MAX_TENSORS {
|
||||
return Err(ProbeError::Malformed("too many outputs"));
|
||||
}
|
||||
info.outputs.push(value_info(w, e)?);
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn opset<R: Read + Seek>(w: &mut Walker<R>, end: u64) -> Result<Option<u64>> {
|
||||
let (mut domain, mut version) = (String::new(), None);
|
||||
while let Some((n, f)) = w.next(end)? {
|
||||
match (n, f) {
|
||||
(1, Wire::Len { start, end: e }) => {
|
||||
domain = w.string(start, e)?;
|
||||
w.seek_to(e)?;
|
||||
}
|
||||
(2, Wire::Varint(v)) => version = Some(v),
|
||||
(_, Wire::Len { end: e, .. }) => w.seek_to(e)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(if domain.is_empty() || domain == "ai.onnx" { version } else { None })
|
||||
}
|
||||
|
||||
/// Read an ONNX model's graph inputs and outputs from any seekable source.
|
||||
pub fn probe_reader<R: Read + Seek>(r: &mut R) -> Result<OnnxInfo> {
|
||||
let total = r.seek(SeekFrom::End(0))?;
|
||||
r.seek(SeekFrom::Start(0))?;
|
||||
if total == 0 {
|
||||
return Err(ProbeError::Malformed("the file is empty"));
|
||||
}
|
||||
let mut w = Walker { r, total, fields: 0 };
|
||||
let (mut info, mut weights, mut saw_graph) = (OnnxInfo::default(), HashSet::new(), false);
|
||||
while let Some((n, f)) = w.next(total)? {
|
||||
match (n, f) {
|
||||
(2, Wire::Len { start, end }) => {
|
||||
info.producer = w.string(start, end)?;
|
||||
w.seek_to(end)?;
|
||||
}
|
||||
(7, Wire::Len { end, .. }) => {
|
||||
saw_graph = true;
|
||||
graph(&mut w, end, &mut info, &mut weights)?;
|
||||
w.seek_to(end)?;
|
||||
}
|
||||
(8, Wire::Len { end, .. }) => {
|
||||
if let Some(v) = opset(&mut w, end)? {
|
||||
info.opset = Some(v);
|
||||
}
|
||||
w.seek_to(end)?;
|
||||
}
|
||||
(_, Wire::Len { end, .. }) => w.seek_to(end)?,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if !saw_graph {
|
||||
return Err(ProbeError::Malformed("no graph"));
|
||||
}
|
||||
info.inputs.retain(|t| !weights.contains(&t.name));
|
||||
if info.inputs.len() > MAX_TENSORS {
|
||||
return Err(ProbeError::Malformed("too many inputs"));
|
||||
}
|
||||
if info.inputs.is_empty() || info.outputs.is_empty() {
|
||||
return Err(ProbeError::Malformed("the graph has no inputs or no outputs"));
|
||||
}
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
/// Read an ONNX file's graph inputs and outputs.
|
||||
pub fn probe_path(path: &Path) -> Result<OnnxInfo> {
|
||||
probe_reader(&mut io::BufReader::new(std::fs::File::open(path)?))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod tests {
|
||||
use super::*;
|
||||
use std::io::Cursor;
|
||||
|
||||
pub(crate) use crate::synthetic::embedder_model as model_bytes;
|
||||
use crate::synthetic::{len_field, value_info_bytes, varint};
|
||||
|
||||
#[test]
|
||||
fn reads_inputs_outputs_and_skips_weights() {
|
||||
let info = probe_reader(&mut Cursor::new(model_bytes(512))).unwrap();
|
||||
assert_eq!(info.producer, "pytorch");
|
||||
assert_eq!(info.opset, Some(17));
|
||||
assert_eq!(info.inputs.len(), 1, "the weight listed as an input is left out");
|
||||
assert_eq!(info.inputs[0].name, "data");
|
||||
assert_eq!(info.inputs[0].elem_type, 1);
|
||||
assert_eq!(info.inputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]);
|
||||
assert_eq!(info.outputs[0].shape, vec![Dim::Dynamic("N".into()), Dim::Fixed(512)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn old_style_files_listing_hundreds_of_weights_as_inputs_still_probe() {
|
||||
let mut graph = Vec::new();
|
||||
len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph);
|
||||
for i in 0..500 {
|
||||
let name = format!("weight{i}");
|
||||
len_field(11, &value_info_bytes(&name, 1, &[Ok(4)]), &mut graph);
|
||||
let mut tensor = Vec::new();
|
||||
len_field(8, name.as_bytes(), &mut tensor);
|
||||
len_field(5, &tensor, &mut graph);
|
||||
}
|
||||
len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(128)]), &mut graph);
|
||||
let mut m = Vec::new();
|
||||
len_field(7, &graph, &mut m);
|
||||
let info = probe_reader(&mut Cursor::new(m)).unwrap();
|
||||
assert_eq!(info.inputs.len(), 1);
|
||||
assert_eq!(info.inputs[0].name, "data");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn many_real_inputs_are_refused() {
|
||||
let mut graph = Vec::new();
|
||||
for i in 0..70 {
|
||||
len_field(11, &value_info_bytes(&format!("in{i}"), 1, &[Ok(4)]), &mut graph);
|
||||
}
|
||||
len_field(12, &value_info_bytes("emb", 1, &[Ok(4)]), &mut graph);
|
||||
let mut m = Vec::new();
|
||||
len_field(7, &graph, &mut m);
|
||||
assert!(probe_reader(&mut Cursor::new(m)).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncated_at_every_length_never_panics() {
|
||||
let bytes = model_bytes(128);
|
||||
for n in 0..bytes.len() {
|
||||
let _ = probe_reader(&mut Cursor::new(bytes[..n].to_vec()));
|
||||
}
|
||||
assert!(probe_reader(&mut Cursor::new(bytes)).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hostile_files_are_errors() {
|
||||
let mut lying = Vec::new();
|
||||
varint(7 << 3 | 2, &mut lying);
|
||||
varint(1 << 40, &mut lying); // a graph 1 TB long in a 10-byte file
|
||||
let mut overflow = Vec::new();
|
||||
varint(7 << 3 | 2, &mut overflow);
|
||||
varint(u64::MAX, &mut overflow);
|
||||
let mut long_number = vec![0x08];
|
||||
long_number.extend_from_slice(&[0xff; 11]);
|
||||
for (what, bytes) in [
|
||||
("empty", Vec::new()),
|
||||
("zeros", vec![0u8; 64]),
|
||||
("text", b"this is not an onnx model at all".to_vec()),
|
||||
("lying length", lying),
|
||||
("overflowing length", overflow),
|
||||
("endless number", long_number),
|
||||
("no graph", {
|
||||
let mut m = Vec::new();
|
||||
len_field(2, b"x", &mut m);
|
||||
m
|
||||
}),
|
||||
("graph without io", {
|
||||
let mut m = Vec::new();
|
||||
len_field(7, &[], &mut m);
|
||||
m
|
||||
}),
|
||||
] {
|
||||
assert!(probe_reader(&mut Cursor::new(bytes)).is_err(), "{what}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn too_many_dimensions_or_tensors_are_errors() {
|
||||
let dims: Vec<std::result::Result<u64, &str>> = (0..40).map(Ok).collect();
|
||||
let mut graph = Vec::new();
|
||||
len_field(11, &value_info_bytes("data", 1, &dims), &mut graph);
|
||||
len_field(12, &value_info_bytes("out", 1, &[Ok(1)]), &mut graph);
|
||||
let mut m = Vec::new();
|
||||
len_field(7, &graph, &mut m);
|
||||
assert!(probe_reader(&mut Cursor::new(m)).is_err());
|
||||
}
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn arbitrary_bytes_never_panic(bytes in proptest::collection::vec(proptest::prelude::any::<u8>(), 0..400)) {
|
||||
let _ = probe_reader(&mut Cursor::new(bytes));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_valid_model_with_flipped_bytes_never_panics(flips in proptest::collection::vec((0usize..400, proptest::prelude::any::<u8>()), 1..6)) {
|
||||
let mut bytes = model_bytes(512);
|
||||
for (at, v) in flips {
|
||||
if let Some(b) = bytes.get_mut(at) {
|
||||
*b = v;
|
||||
}
|
||||
}
|
||||
let _ = probe_reader(&mut Cursor::new(bytes));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
//! What to do with a model file the user drops in: recognise it by hash, or look at its shape and propose a
|
||||
//! manifest the user confirms, or say plainly why it cannot be used yet.
|
||||
|
||||
use crate::manifest::{Commercial, InputSpec, Licence, ModelManifest, OutputSpec, Role, Thresholds};
|
||||
use crate::onnx::{Dim, OnnxInfo, TensorInfo};
|
||||
|
||||
/// ONNX element type for float32.
|
||||
const FLOAT32: u32 = 1;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Suggestion {
|
||||
/// A model we know by its SHA-256: its manifest as it is.
|
||||
Known(ModelManifest),
|
||||
/// An unknown model shaped like a face embedder. `assumptions` is what we guessed, in plain words.
|
||||
Draft { manifest: ModelManifest, assumptions: Vec<String> },
|
||||
/// Not usable yet, and why.
|
||||
Unsupported(String),
|
||||
}
|
||||
|
||||
fn fixed(d: &Dim) -> Option<u64> {
|
||||
match d {
|
||||
Dim::Fixed(v) => Some(*v),
|
||||
Dim::Dynamic(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// An id made of the file name and the start of the hash: lowercase letters, digits, `.` `_` `-`.
|
||||
fn make_id(file_name: &str, sha256: &str) -> String {
|
||||
let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s);
|
||||
let clean: String = stem.chars().map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' }).collect();
|
||||
let clean: String = clean.trim_matches('-').chars().take(40).collect();
|
||||
let short: String = sha256.chars().take(8).collect();
|
||||
match (clean.is_empty(), short.is_empty()) {
|
||||
(true, true) => "custom".into(),
|
||||
(true, false) => format!("custom-{short}"),
|
||||
(false, true) => format!("custom-{clean}"),
|
||||
(false, false) => format!("custom-{clean}-{short}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn display_name(file_name: &str) -> String {
|
||||
let stem = file_name.rsplit_once('.').map_or(file_name, |(s, _)| s);
|
||||
let name: String = stem.chars().filter(|c| !c.is_control()).take(60).collect();
|
||||
if name.trim().is_empty() { "Custom face model".into() } else { format!("{} (your model)", name.trim()) }
|
||||
}
|
||||
|
||||
/// The input of a face embedder: float32, `[batch, 3, height, width]`.
|
||||
fn face_input(t: &TensorInfo) -> Option<(u32, u32)> {
|
||||
if t.elem_type != FLOAT32 || t.shape.len() != 4 || t.shape.get(1).and_then(fixed) != Some(3) {
|
||||
return None;
|
||||
}
|
||||
let size = |d: Option<&Dim>| match d.and_then(fixed) {
|
||||
Some(v) if (16..=1024).contains(&v) => u32::try_from(v).ok(),
|
||||
Some(_) => None,
|
||||
None => Some(112), // a dynamic size: the common face-crop size
|
||||
};
|
||||
Some((size(t.shape.get(3))?, size(t.shape.get(2))?))
|
||||
}
|
||||
|
||||
/// Propose what to do with a probed model. `sha256` and `size` describe the file, `file_name` names it.
|
||||
pub fn suggest(info: &OnnxInfo, sha256: &str, size: u64, file_name: &str) -> Suggestion {
|
||||
if let Some(m) = crate::known::lookup(sha256) {
|
||||
return Suggestion::Known(m);
|
||||
}
|
||||
let [input] = info.inputs.as_slice() else {
|
||||
return Suggestion::Unsupported(format!("the model has {} inputs; face models here take one image input", info.inputs.len()));
|
||||
};
|
||||
let Some((width, height)) = face_input(input) else {
|
||||
return Suggestion::Unsupported("its input is not a float image of shape [batch, 3, height, width]".into());
|
||||
};
|
||||
let [output] = info.outputs.as_slice() else {
|
||||
return Suggestion::Unsupported(
|
||||
"it has several outputs, like a face detector; only models that give one vector per face are supported for now".into(),
|
||||
);
|
||||
};
|
||||
if output.elem_type != FLOAT32 || output.shape.len() != 2 {
|
||||
return Suggestion::Unsupported("its output is not one vector per face (shape [batch, dimension])".into());
|
||||
}
|
||||
let dim = match output.shape.get(1).and_then(fixed).map(u32::try_from) {
|
||||
Some(Ok(d)) if (8..=4096).contains(&d) => d,
|
||||
_ => return Suggestion::Unsupported("its output size is not a fixed number between 8 and 4096".into()),
|
||||
};
|
||||
let manifest = ModelManifest {
|
||||
id: make_id(file_name, sha256),
|
||||
name: display_name(file_name),
|
||||
version: "1".into(),
|
||||
role: Role::Embedder,
|
||||
licence: Licence {
|
||||
name: "Unknown".into(),
|
||||
commercial: Commercial::Unknown,
|
||||
url: None,
|
||||
notice: "LightCraft does not know where this model came from or what its licence allows. Use it only if its licence allows what you are doing, and do not redistribute it.".into(),
|
||||
},
|
||||
source: None,
|
||||
sha256: Some(sha256.to_string()).filter(|h| h.len() == 64),
|
||||
size_bytes: Some(size).filter(|s| *s > 0),
|
||||
provenance: "Unknown.".into(),
|
||||
input: InputSpec { width, height, ..InputSpec::default() },
|
||||
output: OutputSpec::Embedding { dim },
|
||||
thresholds: Thresholds::default(),
|
||||
};
|
||||
let assumptions = vec![
|
||||
format!("Faces are aligned crops of {width} × {height} pixels."),
|
||||
"Pixels are RGB, normalised as (x − 127.5) / 127.5: the usual ArcFace convention, which InsightFace models use.".into(),
|
||||
format!("Each face becomes a vector of {dim} numbers; two faces match when their cosine similarity is high."),
|
||||
];
|
||||
Suggestion::Draft { manifest, assumptions }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::manifest::validate;
|
||||
use crate::onnx::probe_reader;
|
||||
use std::io::Cursor;
|
||||
|
||||
fn tensor(elem: u32, shape: &[Dim]) -> TensorInfo {
|
||||
TensorInfo { name: "t".into(), elem_type: elem, shape: shape.to_vec() }
|
||||
}
|
||||
fn dynamic() -> Dim {
|
||||
Dim::Dynamic("N".into())
|
||||
}
|
||||
fn info(input: TensorInfo, output: TensorInfo) -> OnnxInfo {
|
||||
OnnxInfo { producer: "x".into(), opset: Some(17), inputs: vec![input], outputs: vec![output] }
|
||||
}
|
||||
const SHA: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
|
||||
|
||||
#[test]
|
||||
fn a_known_hash_wins() {
|
||||
let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(10), Dim::Fixed(10)]), tensor(1, &[dynamic(), Dim::Fixed(5)]));
|
||||
assert!(matches!(suggest(&i, crate::known::AURAFACE_SHA256, 1, "whatever.onnx"), Suggestion::Known(m) if m.id == "auraface-v1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unknown_embedder_gets_a_draft_that_validates() {
|
||||
let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)]));
|
||||
let Suggestion::Draft { manifest, assumptions } = suggest(&i, SHA, 5000, "My Model (r50).onnx") else { unreachable_in_test() };
|
||||
assert_eq!(validate(&manifest), Ok(()));
|
||||
assert_eq!(manifest.id, "custom-my-model--r50-01234567");
|
||||
assert_eq!(manifest.output, OutputSpec::Embedding { dim: 512 });
|
||||
assert_eq!((manifest.input.width, manifest.input.height), (112, 112));
|
||||
assert_eq!(manifest.licence.commercial, Commercial::Unknown);
|
||||
assert!(assumptions.iter().any(|a| a.contains("127.5")));
|
||||
}
|
||||
|
||||
#[allow(clippy::panic)]
|
||||
fn unreachable_in_test() -> ! {
|
||||
panic!("expected a draft");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dynamic_sizes_default_and_odd_shapes_are_refused() {
|
||||
let ok = info(tensor(1, &[dynamic(), Dim::Fixed(3), dynamic(), dynamic()]), tensor(1, &[dynamic(), Dim::Fixed(128)]));
|
||||
assert!(matches!(suggest(&ok, SHA, 1, "a.onnx"), Suggestion::Draft { manifest, .. } if manifest.input.width == 112));
|
||||
let refused = [
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(1), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // grey
|
||||
info(tensor(7, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // int64 input
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])), // rank 3
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(1_000_000), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(128)])),
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(257), Dim::Fixed(384)])), // tokens
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), dynamic()])),
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(2)])),
|
||||
info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(1_000_000)])),
|
||||
];
|
||||
for r in refused {
|
||||
assert!(matches!(suggest(&r, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "{r:?}");
|
||||
}
|
||||
let mut two = ok.clone();
|
||||
two.outputs.push(two.outputs[0].clone());
|
||||
assert!(matches!(suggest(&two, SHA, 1, "a.onnx"), Suggestion::Unsupported(_)), "several outputs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_names_cannot_break_the_manifest() {
|
||||
let i = info(tensor(1, &[dynamic(), Dim::Fixed(3), Dim::Fixed(112), Dim::Fixed(112)]), tensor(1, &[dynamic(), Dim::Fixed(512)]));
|
||||
for name in ["", ".onnx", "../../etc/passwd", "\u{7}\u{0}x.onnx", &"x".repeat(5000), "ÅÄÖ 🙂.onnx", "UPPER lower.ONNX"] {
|
||||
let Suggestion::Draft { manifest, .. } = suggest(&i, SHA, 1, name) else { unreachable_in_test() };
|
||||
assert_eq!(validate(&manifest), Ok(()), "{name:?}");
|
||||
}
|
||||
// and with no usable hash at all
|
||||
let Suggestion::Draft { manifest, .. } = suggest(&i, "", 0, "x.onnx") else { unreachable_in_test() };
|
||||
assert_eq!(validate(&manifest), Ok(()));
|
||||
assert!(manifest.sha256.is_none() && manifest.size_bytes.is_none());
|
||||
}
|
||||
|
||||
/// Opt-in: `LC_FACE_MODELS=/folder/with/onnx/files cargo test -p lightcraft-faces -- --ignored --nocapture`
|
||||
/// probes real models (never committed) and prints what would be suggested, with timings.
|
||||
#[test]
|
||||
#[ignore = "needs real model files: set LC_FACE_MODELS"]
|
||||
fn probing_real_models() {
|
||||
let Some(dir) = std::env::var_os("LC_FACE_MODELS") else { return };
|
||||
for entry in std::fs::read_dir(dir).unwrap().flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().is_none_or(|e| e != "onnx") {
|
||||
continue;
|
||||
}
|
||||
let t = std::time::Instant::now();
|
||||
let sha = crate::hash::sha256_file(&path).unwrap();
|
||||
let hashed = t.elapsed();
|
||||
let t = std::time::Instant::now();
|
||||
let info = crate::onnx::probe_path(&path);
|
||||
let probed = t.elapsed();
|
||||
let size = std::fs::metadata(&path).unwrap().len();
|
||||
let name = path.file_name().unwrap().to_string_lossy().into_owned();
|
||||
let verdict = match &info {
|
||||
Ok(i) => match suggest(i, &sha, size, &name) {
|
||||
Suggestion::Known(m) => format!("KNOWN {} ({:?})", m.id, m.licence.commercial),
|
||||
Suggestion::Draft { manifest, .. } => format!("DRAFT {} {:?}", manifest.id, manifest.output),
|
||||
Suggestion::Unsupported(why) => format!("UNSUPPORTED: {why}"),
|
||||
},
|
||||
Err(e) => format!("PROBE ERROR: {e}"),
|
||||
};
|
||||
println!("{name}: {} MB | hash {hashed:?} | probe {probed:?} | {verdict}", size / 1_048_576);
|
||||
if let Ok(i) = info {
|
||||
println!(
|
||||
" in {:?}
|
||||
out {:?}",
|
||||
i.inputs, i.outputs
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn probing_then_suggesting_end_to_end() {
|
||||
let probed = probe_reader(&mut Cursor::new(crate::onnx::tests::model_bytes(512))).unwrap();
|
||||
assert!(
|
||||
matches!(suggest(&probed, SHA, 9, "m.onnx"), Suggestion::Draft { manifest, .. } if manifest.output == OutputSpec::Embedding { dim: 512 })
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
//! A minimal, valid ONNX model made in memory: for tests, and later for the install self-check. It is not a
|
||||
//! working network (it has no layers), only a graph with the right inputs, outputs and a weight blob.
|
||||
|
||||
pub(crate) fn varint(mut v: u64, out: &mut Vec<u8>) {
|
||||
loop {
|
||||
let b = (v & 0x7f) as u8;
|
||||
v >>= 7;
|
||||
if v == 0 {
|
||||
out.push(b);
|
||||
return;
|
||||
}
|
||||
out.push(b | 0x80);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn len_field(n: u64, body: &[u8], out: &mut Vec<u8>) {
|
||||
varint(n << 3 | 2, out);
|
||||
varint(body.len() as u64, out);
|
||||
out.extend_from_slice(body);
|
||||
}
|
||||
|
||||
pub(crate) fn varint_field(n: u64, v: u64, out: &mut Vec<u8>) {
|
||||
varint(n << 3, out);
|
||||
varint(v, out);
|
||||
}
|
||||
|
||||
/// A graph input or output: `Ok(n)` is a fixed dimension, `Err(name)` a named (dynamic) one.
|
||||
pub fn value_info_bytes(name: &str, elem: u64, dims: &[std::result::Result<u64, &str>]) -> Vec<u8> {
|
||||
let mut shape = Vec::new();
|
||||
for d in dims {
|
||||
let mut dim = Vec::new();
|
||||
match d {
|
||||
Ok(v) => varint_field(1, *v, &mut dim),
|
||||
Err(p) => len_field(2, p.as_bytes(), &mut dim),
|
||||
}
|
||||
len_field(1, &dim, &mut shape);
|
||||
}
|
||||
let mut tensor = Vec::new();
|
||||
varint_field(1, elem, &mut tensor);
|
||||
len_field(2, &shape, &mut tensor);
|
||||
let mut ty = Vec::new();
|
||||
len_field(1, &tensor, &mut ty);
|
||||
let mut vi = Vec::new();
|
||||
len_field(1, name.as_bytes(), &mut vi);
|
||||
len_field(2, &ty, &mut vi);
|
||||
vi
|
||||
}
|
||||
|
||||
/// A small valid model shaped like a face embedder: float input `data` `[N, 3, 112, 112]`, an old-style
|
||||
/// weight also listed as an input, a weight with bulk data, float output `emb` `[N, output_dim]`, opset 17.
|
||||
pub fn embedder_model(output_dim: u64) -> Vec<u8> {
|
||||
let mut graph = Vec::new();
|
||||
len_field(11, &value_info_bytes("data", 1, &[Err("N"), Ok(3), Ok(112), Ok(112)]), &mut graph);
|
||||
len_field(11, &value_info_bytes("w", 1, &[Ok(4)]), &mut graph);
|
||||
let mut tensor = Vec::new();
|
||||
varint_field(1, 4, &mut tensor);
|
||||
len_field(8, b"w", &mut tensor);
|
||||
len_field(9, &[7u8; 5000], &mut tensor);
|
||||
len_field(5, &tensor, &mut graph);
|
||||
len_field(12, &value_info_bytes("emb", 1, &[Err("N"), Ok(output_dim)]), &mut graph);
|
||||
let mut ops = Vec::new();
|
||||
len_field(1, b"", &mut ops);
|
||||
varint_field(2, 17, &mut ops);
|
||||
let mut model = Vec::new();
|
||||
varint_field(1, 8, &mut model);
|
||||
len_field(2, b"pytorch", &mut model);
|
||||
len_field(7, &graph, &mut model);
|
||||
len_field(8, &ops, &mut model);
|
||||
model
|
||||
}
|
||||
@@ -31,7 +31,7 @@ impl Drop for Headless {
|
||||
|
||||
impl Default for Headless {
|
||||
fn default() -> Self {
|
||||
Self::new(Session::new().with_fs())
|
||||
Self::new(Session::new().with_fs().with_default_face_models())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,7 +42,7 @@ impl Headless {
|
||||
|
||||
/// A headless session with the procedurally generated demo library.
|
||||
pub fn demo() -> Self {
|
||||
Self::new(Session::with_demo().with_fs())
|
||||
Self::new(Session::with_demo().with_fs().with_default_face_models())
|
||||
}
|
||||
|
||||
fn photo_or_active(&self, p: &Value) -> Result<PhotoId, String> {
|
||||
|
||||
@@ -26,6 +26,7 @@ log = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
lightcraft-codecs = { workspace = true }
|
||||
lightcraft-faces = { workspace = true }
|
||||
|
||||
[build-dependencies]
|
||||
serde_json = { workspace = true }
|
||||
|
||||
@@ -620,6 +620,60 @@ mod tests {
|
||||
h.settle(SETTLE);
|
||||
}
|
||||
|
||||
/// Adding a face model: the dialog shows the file's terms, the model is installed only once they are
|
||||
/// accepted, and a file LightCraft cannot use only says why.
|
||||
#[test]
|
||||
fn adding_a_face_model_shows_its_terms_and_installs_only_once_accepted() {
|
||||
use crate::state::Dialog;
|
||||
let mut h = demo([1400.0, 900.0]);
|
||||
let t = Duration::from_secs(10);
|
||||
let dir = std::env::temp_dir().join(format!("lc-ui-facemodel-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
h.app.session.face_models_dir = Some(dir.join("models"));
|
||||
let model = dir.join("Mine.onnx");
|
||||
std::fs::write(&model, lightcraft_faces::synthetic::embedder_model(512)).unwrap();
|
||||
let open = |h: &mut Headless, path: &std::path::Path| {
|
||||
h.request("engine.execute", json!({"command": "dialog.faceModel", "params": {"path": path.to_string_lossy()}}), t)
|
||||
};
|
||||
|
||||
// the dialog opens for the chosen file, with nothing accepted and nothing installed
|
||||
assert_eq!(open(&mut h, &model)["ok"], true);
|
||||
h.settle(SETTLE);
|
||||
h.step();
|
||||
let Some(dlg) = h.app.ui.dialog.clone() else { panic!("no dialog") };
|
||||
let Dialog::FaceModel { info, accepted, .. } = &dlg else { panic!("{dlg:?}") };
|
||||
assert!(!accepted);
|
||||
assert_eq!(info["kind"], "draft");
|
||||
assert!(!dir.join("models").exists());
|
||||
// OK does nothing yet
|
||||
assert!(crate::panels::dialogs::confirm_dialog(&mut h.app, &dlg).is_err());
|
||||
assert!(!dir.join("models").exists());
|
||||
// ticking the box and confirming installs it
|
||||
let r = h.request("ui.clickWidget", json!({"id": "check:faceModel.accept"}), t);
|
||||
assert_eq!(r["ok"], true, "{r}");
|
||||
h.step();
|
||||
assert!(matches!(&h.app.ui.dialog, Some(Dialog::FaceModel { accepted: true, .. })));
|
||||
let r = h.request("ui.dialog.confirm", json!({}), t);
|
||||
assert_eq!(r["ok"], true, "{r}");
|
||||
let listed = h.app.run("faces.models.list", json!({})).unwrap();
|
||||
assert!(listed["models"].as_array().unwrap().iter().any(|m| m["installed"] == true && m["known"] == false), "{listed}");
|
||||
assert!(h.app.ui.dialog.is_none());
|
||||
|
||||
// a file that is not a usable model says why and offers no install
|
||||
let junk = dir.join("junk.onnx");
|
||||
std::fs::write(&junk, b"not a model").unwrap();
|
||||
assert_eq!(open(&mut h, &junk)["ok"], true);
|
||||
h.step();
|
||||
let Some(Dialog::FaceModel { info, .. }) = h.app.ui.dialog.clone() else { panic!("no dialog") };
|
||||
assert_eq!(info["kind"], "unsupported");
|
||||
// a path that does not exist is an error, not a dialog
|
||||
h.app.ui.dialog = None;
|
||||
assert_eq!(open(&mut h, &dir.join("nope.onnx"))["ok"], false);
|
||||
assert!(h.app.ui.dialog.is_none());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
/// Profile browser: live variant thumbnails, hover previews in the loupe without touching the
|
||||
/// photo or its history, click applies, the star toggles the favourite.
|
||||
#[test]
|
||||
|
||||
@@ -106,6 +106,8 @@ pub struct Services {
|
||||
pub picker: Option<pick::Picker>,
|
||||
/// Show an open dialog for photos; returns paths.
|
||||
pub pick_files: Option<PickFiles>,
|
||||
/// Open dialog for a face model (`.onnx`; Settings ▸ Faces ▸ Add a model file…).
|
||||
pub pick_model_file: Option<PickFiles>,
|
||||
/// Open dialog for preset files (`.lcpreset`, `.xmp`, `.lrtemplate`, `.zip`, `.dng`, Luminar `.lmp` / `.mplumpack`).
|
||||
pub pick_preset_files: Option<PickFiles>,
|
||||
/// Open dialog for a GPS track log (`.gpx`; Photo ▸ Auto-Tag from Tracklog…).
|
||||
@@ -764,6 +766,13 @@ impl LightcraftApp {
|
||||
let dropped: Vec<String> =
|
||||
ctx.input(|i| i.raw.dropped_files.iter().map(|f| f.path().to_string_lossy().to_string()).filter(|p| !p.is_empty()).collect());
|
||||
// preset files import as presets, everything else as photos
|
||||
// a dropped face model opens its licence dialog; presets and photos as before
|
||||
let (models, dropped): (Vec<String>, Vec<String>) = dropped.into_iter().partition(|p| is_model_file(p));
|
||||
if let Some(model) = models.first()
|
||||
&& let Err(e) = panels::faces::open_dialog(self, model)
|
||||
{
|
||||
self.toast(ctx, e);
|
||||
}
|
||||
let (presets, photos): (Vec<String>, Vec<String>) = dropped.into_iter().partition(|p| is_preset_file(p));
|
||||
if !presets.is_empty() {
|
||||
let _ = self.run("file.importPresets", serde_json::json!({"paths": presets}));
|
||||
@@ -959,6 +968,11 @@ pub fn is_bw(d: &lightcraft_develop::DevelopSettings) -> bool {
|
||||
d.treatment == lightcraft_develop::Treatment::Bw || d.profile.id == "lc.mono" || d.profile.id.starts_with("lc.bw.")
|
||||
}
|
||||
|
||||
/// A dropped face model file (`.onnx`).
|
||||
pub fn is_model_file(path: &str) -> bool {
|
||||
std::path::Path::new(path).extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx"))
|
||||
}
|
||||
|
||||
/// Files dropped on the window that are presets rather than photos.
|
||||
pub fn is_preset_file(path: &str) -> bool {
|
||||
let ext = std::path::Path::new(path).extension().map(|e| e.to_string_lossy().to_ascii_lowercase()).unwrap_or_default();
|
||||
@@ -967,6 +981,16 @@ pub fn is_preset_file(path: &str) -> bool {
|
||||
|
||||
#[cfg(test)]
|
||||
mod drop_tests {
|
||||
#[test]
|
||||
fn dropped_face_models_are_told_apart() {
|
||||
for p in ["/a/model.onnx", "/a/dir/M.ONNX"] {
|
||||
assert!(super::is_model_file(p), "{p}");
|
||||
}
|
||||
for p in ["/a/model.onnx.jpg", "/a/onnx", "/a/b.png"] {
|
||||
assert!(!super::is_model_file(p), "{p}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dropped_presets_are_told_apart_from_photos() {
|
||||
for p in ["/a/Look.lrtemplate", "/a/b.XMP", "/a/pack.zip", "/a/x.lcpreset", "/a/Magic Hour.mplumpack", "/a/Pop.lmp", "/a/Bundle.LMP"] {
|
||||
@@ -991,6 +1015,8 @@ pub struct Caches {
|
||||
album_counts: Option<(u64, std::sync::Arc<std::collections::HashMap<lightcraft_catalog::AlbumId, usize>>)>,
|
||||
/// How often the album counts were recomputed (tests check that unchanged frames don't).
|
||||
pub album_count_scans: usize,
|
||||
/// Bumped when a face model is installed, removed or chosen, so Settings re-reads the list at once.
|
||||
pub faces_epoch: u64,
|
||||
/// The grid's date runs, layout and indexes (by the visible list's generation).
|
||||
pub grid: panels::grid::GridCache,
|
||||
/// What the grid did on its frames (benchmarks and tests check unchanged frames stay cheap).
|
||||
|
||||
@@ -135,6 +135,7 @@ pub const UI_COMMANDS: &[UiCommand] = &[
|
||||
("tool.keywordPainter", "Keyword Painter", None, ""),
|
||||
("view.gridInfo", "Grid Info", None, ""),
|
||||
("dialog.allMetadata", "All Metadata…", None, "Photo"),
|
||||
("dialog.faceModel", "Add Face Model…", None, ""),
|
||||
("dialog.newSmartAlbum", "New Smart Album from Filter…", Some("Cmd+Alt+N"), "File"),
|
||||
("dialog.createPreset", "Create Preset…", Some("Cmd+Shift+P"), "Photo"),
|
||||
("dialog.autoStack", "Auto-Stack by Capture Time…", None, "Photo>Stack"),
|
||||
@@ -433,7 +434,7 @@ pub fn run_ui_command(app: &mut LightcraftApp, id: &str, p: &Value) -> Option<Re
|
||||
"app.settings" => {
|
||||
let tab = p.get("tab").and_then(Value::as_str).unwrap_or("general");
|
||||
if !crate::panels::settings::TABS.iter().any(|(id, _)| *id == tab) {
|
||||
return Some(Err(format!("unknown settings tab `{tab}` (general|import|performance|interface)")));
|
||||
return Some(Err(format!("unknown settings tab `{tab}` (general|import|performance|interface|faces)")));
|
||||
}
|
||||
app.ui.dialog = Some(Dialog::Settings { tab: tab.into() });
|
||||
Ok(Value::Null)
|
||||
@@ -906,6 +907,16 @@ pub fn run_ui_command(app: &mut LightcraftApp, id: &str, p: &Value) -> Option<Re
|
||||
"merge.panoramaLast" => crate::merge::start_last(app, "merge.panorama"),
|
||||
"merge.hdrPanoramaLast" => crate::merge::start_last(app, "merge.hdrPanorama"),
|
||||
"dialog.mergeHdr" => crate::merge::open(app, "merge.hdr"),
|
||||
"dialog.faceModel" => {
|
||||
let path = match p.get("path").and_then(Value::as_str) {
|
||||
Some(x) => Some(x.to_string()),
|
||||
None => app.services.pick_model_file.as_mut().and_then(|f| f().into_iter().next()),
|
||||
};
|
||||
match path {
|
||||
Some(path) => crate::panels::faces::open_dialog(app, &path),
|
||||
None => Ok(Value::Null),
|
||||
}
|
||||
}
|
||||
"dialog.mergePanorama" => crate::merge::open(app, "merge.panorama"),
|
||||
"dialog.mergeHdrPanorama" => crate::merge::open(app, "merge.hdrPanorama"),
|
||||
"app.about" => {
|
||||
|
||||
@@ -85,6 +85,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
Dialog::NewSmartAlbum { .. } => "Create Smart Album",
|
||||
Dialog::AllMetadata { .. } => "All Metadata",
|
||||
Dialog::SystemInfo { .. } => "System Info",
|
||||
Dialog::FaceModel { .. } => "Add Face Model",
|
||||
Dialog::WhatsNew => "What's New",
|
||||
Dialog::Cull { .. } => "Assisted Culling",
|
||||
Dialog::SmartRules { id: None, .. } => "New Smart Album",
|
||||
@@ -123,6 +124,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
Dialog::Import { .. } => crate::import::DIALOG_SIZE[0],
|
||||
Dialog::SmartRules { .. } => 680.0,
|
||||
Dialog::AllMetadata { .. } => 620.0,
|
||||
Dialog::FaceModel { .. } => 460.0,
|
||||
_ => 380.0,
|
||||
})
|
||||
.show(ctx, |ui| {
|
||||
@@ -174,6 +176,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
}
|
||||
});
|
||||
}
|
||||
Dialog::FaceModel { info, accepted, .. } => crate::panels::faces::model_dialog(app, ui, &t, info, accepted),
|
||||
Dialog::SystemInfo { rows } => {
|
||||
egui::Grid::new("sysinfo").num_columns(2).spacing([16.0, 4.0]).striped(true).show(ui, |ui| {
|
||||
for (k, v) in rows.iter() {
|
||||
@@ -784,7 +787,9 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
}
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| {
|
||||
let informational = matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. });
|
||||
// a model file LightCraft cannot use has nothing to confirm
|
||||
let unusable_model = matches!(&dlg, Dialog::FaceModel { info, .. } if info["kind"] == "unsupported");
|
||||
let informational = unusable_model || matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. });
|
||||
let sam = &app.session.segmenter;
|
||||
let (sam_installed, sam_running, sam_failed) = (sam.installed(), sam.download_status().running, sam.download_status().error.is_some());
|
||||
// no download location in this build: nothing to offer but the manual install
|
||||
@@ -810,6 +815,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
add_label.as_str()
|
||||
}
|
||||
Dialog::Merge { .. } => "Merge",
|
||||
Dialog::FaceModel { .. } if !informational => "Install",
|
||||
Dialog::ConfirmDelete { .. } => "Delete",
|
||||
Dialog::RemoveFolder { .. } => "Remove",
|
||||
Dialog::SamModel { then: Some(_), .. } if sam_installed => "Continue",
|
||||
@@ -819,7 +825,9 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
||||
_ if informational => "Close",
|
||||
_ => "OK",
|
||||
};
|
||||
let r = (!ok.is_empty()).then(|| ui.button(crate::i18n::tr(ok)));
|
||||
// installing waits for the licence to be accepted
|
||||
let can_confirm = informational || !matches!(&dlg, Dialog::FaceModel { accepted: false, .. });
|
||||
let r = (!ok.is_empty()).then(|| ui.add_enabled(can_confirm, egui::Button::new(crate::i18n::tr(ok))));
|
||||
if let Some(r) = &r {
|
||||
crate::widgets::register(ui.ctx(), "button:dialogOk", r.rect);
|
||||
}
|
||||
@@ -1025,6 +1033,7 @@ pub fn confirm_dialog(app: &mut LightcraftApp, dlg: &Dialog) -> Result<serde_jso
|
||||
Dialog::RenameKeyword { from, to } => app.run("keyword.rename", json!({"from": from, "to": to})),
|
||||
Dialog::MergeKeywords { from, into } => app.run("keyword.merge", json!({"from": from, "into": into})),
|
||||
Dialog::AutoStack { gap } => app.run("stack.auto", json!({"gap": gap})),
|
||||
Dialog::FaceModel { path, accepted, .. } => crate::panels::faces::install(app, path, *accepted),
|
||||
Dialog::AllMetadata { .. } | Dialog::SystemInfo { .. } | Dialog::WhatsNew => Ok(serde_json::Value::Null),
|
||||
Dialog::Cull { reject_below, pick_best } => {
|
||||
let mut p = json!({"pickBest": pick_best});
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
//! Faces: the Settings ▸ Faces tab (the on/off toggle and the list of face models) and the dialog that
|
||||
//! shows a model file's licence terms before it is installed.
|
||||
//!
|
||||
//! Adding a model is: pick or drop a `.onnx` file → the engine inspects it (`faces.models.inspect`) →
|
||||
//! this dialog shows what it is and its terms → the user ticks "I accept" → `faces.models.install`.
|
||||
//! LightCraft never fetches a model by itself: "Get…" opens the model's own page in the browser.
|
||||
|
||||
use egui::RichText;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::settings::{check, heading, hint};
|
||||
use crate::LightcraftApp;
|
||||
use crate::theme::Tokens;
|
||||
use crate::widgets::register;
|
||||
|
||||
/// How long a read of the model list is reused (it is a few small files, but not for every frame).
|
||||
const REFRESH_SECS: f64 = 1.5;
|
||||
|
||||
fn list_id() -> egui::Id {
|
||||
egui::Id::new("faces-model-list")
|
||||
}
|
||||
|
||||
/// The engine's model list, re-read at most every [`REFRESH_SECS`], or at once after an action changed it.
|
||||
fn models(app: &mut LightcraftApp, ui: &egui::Ui) -> Value {
|
||||
let now = ui.input(|i| i.time);
|
||||
let epoch = app.caches.faces_epoch;
|
||||
if let Some((e, at, v)) = ui.data(|d| d.get_temp::<(u64, f64, Value)>(list_id()))
|
||||
&& e == epoch
|
||||
&& now - at < REFRESH_SECS
|
||||
{
|
||||
return v;
|
||||
}
|
||||
let v = app.run("faces.models.list", json!({})).unwrap_or(Value::Null);
|
||||
ui.data_mut(|d| d.insert_temp(list_id(), (epoch, now, v.clone())));
|
||||
v
|
||||
}
|
||||
|
||||
/// A download-style size: decimal megabytes, as file managers and model pages show them.
|
||||
fn mb(bytes: Option<u64>) -> String {
|
||||
match bytes {
|
||||
Some(b) if b >= 10_000_000 => format!("{} MB", (b as f64 / 1e6).round() as u64),
|
||||
Some(b) if b >= 1_000_000 => format!("{:.1} MB", b as f64 / 1e6),
|
||||
Some(b) => format!("{} KB", (b as f64 / 1e3).round().max(1.0) as u64),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// "its input is not an image" → "Its input is not an image."
|
||||
fn sentence(s: &str) -> String {
|
||||
let s = s.trim();
|
||||
let mut c = s.chars();
|
||||
let mut out: String = c.next().map(|f| f.to_uppercase().collect()).unwrap_or_default();
|
||||
out.push_str(c.as_str());
|
||||
if !out.is_empty() && !out.ends_with(['.', '!', '?']) {
|
||||
out.push('.');
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// "Apache-2.0 · commercial use allowed"
|
||||
fn licence_line(m: &Value) -> String {
|
||||
let name = m["licence"]["name"].as_str().filter(|s| !s.is_empty()).unwrap_or("Unknown licence");
|
||||
let terms = match m["licence"]["commercial"].as_str() {
|
||||
Some("yes") => "commercial use allowed",
|
||||
Some("no") => "non-commercial use only",
|
||||
_ => "terms unclear",
|
||||
};
|
||||
format!("{name} · {terms}")
|
||||
}
|
||||
|
||||
fn open_page(app: &mut LightcraftApp, url: &str) {
|
||||
if let Some(f) = app.services.open_url.as_mut() {
|
||||
let _ = f(url);
|
||||
}
|
||||
}
|
||||
|
||||
/// The Settings ▸ Faces tab.
|
||||
pub fn settings_tab(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens) {
|
||||
let list = models(app, ui);
|
||||
heading(ui, t, "Face recognition");
|
||||
if list["dir"].is_null() {
|
||||
hint(ui, t, "This build has nowhere to keep face models, so the model list is not available. (The desktop app has.)");
|
||||
return;
|
||||
}
|
||||
let mut on = list["enabled"].as_bool().unwrap_or(false);
|
||||
if check(ui, "settings.facesEnabled", &mut on, "Recognise faces (experimental)") {
|
||||
let _ = app.run("faces.enable", json!({"enabled": on}));
|
||||
app.caches.faces_epoch += 1;
|
||||
}
|
||||
hint(
|
||||
ui,
|
||||
t,
|
||||
"Off by default. LightCraft already reads the face names other apps wrote into your photos; this is for models that suggest who is in a photo. Everything stays on your computer.",
|
||||
);
|
||||
if list["runtime"].as_bool() != Some(true) {
|
||||
hint(ui, t, "Models can be added and chosen now. Running them arrives in a later update.");
|
||||
}
|
||||
let all: Vec<Value> = list["models"].as_array().cloned().unwrap_or_default();
|
||||
heading(ui, t, "Face detector");
|
||||
for m in all.iter().filter(|m| m["role"] == "detector") {
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(m["name"].as_str().unwrap_or("")).color(t.text));
|
||||
ui.label(RichText::new(format!("{} · included", mb(m["sizeBytes"].as_u64()))).color(t.text_dim));
|
||||
});
|
||||
}
|
||||
heading(ui, t, "Recognition models");
|
||||
for m in all.iter().filter(|m| m["role"] == "embedder") {
|
||||
model_row(app, ui, t, m);
|
||||
}
|
||||
ui.add_space(6.0);
|
||||
ui.horizontal(|ui| {
|
||||
let can = app.services.pick_model_file.is_some();
|
||||
let r = ui.add_enabled(can, egui::Button::new("Add a model file…"));
|
||||
register(ui.ctx(), "faces:addModel", r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("dialog.faceModel", json!({}));
|
||||
}
|
||||
ui.label(RichText::new("or drop a .onnx file on the window").color(t.text_dim));
|
||||
});
|
||||
}
|
||||
|
||||
fn model_row(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens, m: &Value) {
|
||||
let id = m["id"].as_str().unwrap_or("").to_string();
|
||||
let (installed, selected) = (m["installed"].as_bool() == Some(true), m["selected"].as_bool() == Some(true));
|
||||
egui::Frame::NONE.inner_margin(egui::Margin::symmetric(0, 3)).show(ui, |ui| {
|
||||
ui.horizontal(|ui| {
|
||||
ui.vertical(|ui| {
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(m["name"].as_str().unwrap_or("")).color(t.text));
|
||||
let size = mb(m["sizeBytes"].as_u64());
|
||||
if !size.is_empty() {
|
||||
ui.label(RichText::new(size).color(t.text_dim));
|
||||
}
|
||||
if selected {
|
||||
ui.label(RichText::new("in use").color(t.accent));
|
||||
} else if installed {
|
||||
ui.label(RichText::new("installed").color(t.text_dim));
|
||||
}
|
||||
});
|
||||
ui.label(RichText::new(licence_line(m)).font(t.font(11.5)).color(if m["licence"]["commercial"] == "yes" {
|
||||
t.text_dim
|
||||
} else {
|
||||
t.caution
|
||||
}));
|
||||
});
|
||||
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
|
||||
if installed {
|
||||
let r = ui.button("Remove");
|
||||
register(ui.ctx(), format!("faces:remove:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("faces.models.remove", json!({"id": id}));
|
||||
app.caches.faces_epoch += 1;
|
||||
}
|
||||
if !selected {
|
||||
let r = ui.button("Use");
|
||||
register(ui.ctx(), format!("faces:use:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
let _ = app.run("faces.models.select", json!({"id": id}));
|
||||
app.caches.faces_epoch += 1;
|
||||
}
|
||||
}
|
||||
} else if let Some(url) = m["source"].as_str() {
|
||||
let r = ui
|
||||
.button("Get…")
|
||||
.on_hover_text("Opens the model's own page in your browser. Download the .onnx file there, then add it here.");
|
||||
register(ui.ctx(), format!("faces:get:{id}"), r.rect);
|
||||
if r.clicked() {
|
||||
open_page(app, url);
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/// The body of the "Add Face Model" dialog: what the file is, its terms, and the accept box.
|
||||
/// `info` is the engine's `faces.models.inspect` answer.
|
||||
pub fn model_dialog(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens, info: &Value, accepted: &mut bool) {
|
||||
let kind = info["kind"].as_str().unwrap_or("unsupported");
|
||||
let file = info["fileName"].as_str().unwrap_or("");
|
||||
if kind == "unsupported" {
|
||||
ui.label(RichText::new("This file cannot be used yet").font(t.semibold(13.5)).color(t.caution));
|
||||
ui.add(
|
||||
egui::Label::new(
|
||||
RichText::new(sentence(info["reason"].as_str().unwrap_or("It is not a face recognition model LightCraft understands")))
|
||||
.color(t.text_label),
|
||||
)
|
||||
.wrap(),
|
||||
);
|
||||
ui.label(RichText::new(format!("{file} · {}", mb(info["sizeBytes"].as_u64()))).color(t.text_dim));
|
||||
return;
|
||||
}
|
||||
let m = &info["model"];
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(m["name"].as_str().unwrap_or(file)).font(t.semibold(14.0)).color(t.text));
|
||||
ui.label(RichText::new(mb(info["sizeBytes"].as_u64())).color(t.text_dim));
|
||||
});
|
||||
if info["alreadyInstalled"] == true {
|
||||
ui.label(RichText::new("Already installed. Installing again is harmless.").color(t.text_dim));
|
||||
}
|
||||
let commercial = m["licence"]["commercial"].as_str().unwrap_or("unknown");
|
||||
ui.label(RichText::new(licence_line(m)).font(t.semibold(12.5)).color(if commercial == "yes" { t.text } else { t.caution }));
|
||||
let notice = m["licence"]["notice"].as_str().unwrap_or("");
|
||||
if !notice.is_empty() {
|
||||
ui.add(egui::Label::new(RichText::new(notice).color(t.text_label)).wrap());
|
||||
}
|
||||
let provenance = m["provenance"].as_str().unwrap_or("");
|
||||
if !provenance.is_empty() {
|
||||
ui.add(egui::Label::new(RichText::new(format!("Trained on: {provenance}")).color(t.text_dim)).wrap());
|
||||
}
|
||||
if kind == "draft" {
|
||||
ui.add_space(2.0);
|
||||
ui.label(RichText::new("LightCraft does not know this model, so it assumed:").color(t.text_label));
|
||||
for a in info["assumptions"].as_array().into_iter().flatten().filter_map(Value::as_str) {
|
||||
ui.add(egui::Label::new(RichText::new(format!("• {a}")).font(t.font(12.0)).color(t.text_dim)).wrap());
|
||||
}
|
||||
}
|
||||
if let Some(url) = m["source"].as_str().or(m["licence"]["url"].as_str()) {
|
||||
let r = ui.link("Open the model's page");
|
||||
register(ui.ctx(), "faceModel:page", r.rect);
|
||||
if r.clicked() {
|
||||
open_page(app, url);
|
||||
}
|
||||
}
|
||||
ui.add_space(4.0);
|
||||
check(ui, "faceModel.accept", accepted, "I have read these terms and accept them for my own use");
|
||||
ui.label(RichText::new("The model is kept on this computer only. LightCraft never uploads or shares it.").font(t.font(11.5)).color(t.text_dim));
|
||||
}
|
||||
|
||||
/// Inspect `path` and open the dialog for it (also what dropping a `.onnx` file on the window does).
|
||||
pub fn open_dialog(app: &mut LightcraftApp, path: &str) -> Result<Value, String> {
|
||||
let info = app.run("faces.models.inspect", json!({"path": path}))?;
|
||||
app.ui.dialog = Some(crate::state::Dialog::FaceModel { path: path.to_string(), info, accepted: false });
|
||||
Ok(Value::Null)
|
||||
}
|
||||
|
||||
/// The dialog's OK: install the model (the engine refuses without the acceptance).
|
||||
pub fn install(app: &mut LightcraftApp, path: &str, accepted: bool) -> Result<Value, String> {
|
||||
if !accepted {
|
||||
return Err("Tick the box to accept the model's terms first".into());
|
||||
}
|
||||
let r = app.run("faces.models.install", json!({"path": path, "acknowledged": true}));
|
||||
app.caches.faces_epoch += 1;
|
||||
r
|
||||
}
|
||||
@@ -7,6 +7,7 @@ pub mod crop_overlay;
|
||||
pub mod detail;
|
||||
pub mod dialogs;
|
||||
pub mod edit;
|
||||
pub mod faces;
|
||||
pub mod filterbar;
|
||||
pub mod grid;
|
||||
pub mod keymap;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! The Settings dialog (⌘,): General, Import, Performance, Interface.
|
||||
//! The Settings dialog (⌘,): General, Import, Performance, Interface, Faces.
|
||||
//!
|
||||
//! Changes apply immediately (no OK/Cancel). Where they are stored:
|
||||
//! - **app settings** ([`crate::state::AppSettings`]: startup view, delete confirmation, GPU,
|
||||
@@ -17,7 +17,8 @@ use crate::theme::Tokens;
|
||||
use crate::widgets::register;
|
||||
|
||||
/// (id, label) of the tabs, in order.
|
||||
pub const TABS: &[(&str, &str)] = &[("general", "General"), ("import", "Import"), ("performance", "Performance"), ("interface", "Interface")];
|
||||
pub const TABS: &[(&str, &str)] =
|
||||
&[("general", "General"), ("import", "Import"), ("performance", "Performance"), ("interface", "Interface"), ("faces", "Faces")];
|
||||
|
||||
/// Thumbnail cache sizes offered (MB).
|
||||
const CACHE_SIZES: [u32; 5] = [512, 1024, 2048, 4096, 8192];
|
||||
@@ -43,11 +44,12 @@ pub fn body(app: &mut LightcraftApp, ui: &mut egui::Ui, tab: &mut String) {
|
||||
"import" => import_tab(app, ui, &t),
|
||||
"performance" => performance_tab(app, ui, &t),
|
||||
"interface" => interface_tab(app, ui, &t),
|
||||
"faces" => super::faces::settings_tab(app, ui, &t),
|
||||
_ => general_tab(app, ui, &t),
|
||||
}
|
||||
}
|
||||
|
||||
fn heading(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
pub(super) fn heading(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
ui.add_space(4.0);
|
||||
ui.label(RichText::new(crate::i18n::tr(text)).font(t.semibold(12.5)).color(t.text));
|
||||
}
|
||||
@@ -63,12 +65,12 @@ fn row<R>(ui: &mut egui::Ui, t: &Tokens, label: &str, add: impl FnOnce(&mut egui
|
||||
.inner
|
||||
}
|
||||
|
||||
fn hint(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
pub(super) fn hint(ui: &mut egui::Ui, t: &Tokens, text: &str) {
|
||||
ui.label(RichText::new(crate::i18n::tr(text)).size(11.0).color(t.text_dim));
|
||||
}
|
||||
|
||||
/// A checkbox addressable as `check:{id}`; true when toggled.
|
||||
fn check(ui: &mut egui::Ui, id: &str, value: &mut bool, label: &str) -> bool {
|
||||
pub(super) fn check(ui: &mut egui::Ui, id: &str, value: &mut bool, label: &str) -> bool {
|
||||
let r = ui.checkbox(value, crate::i18n::tr(label));
|
||||
register(ui.ctx(), format!("check:{id}"), r.rect);
|
||||
r.changed()
|
||||
|
||||
@@ -535,6 +535,13 @@ pub enum Dialog {
|
||||
pick_best: bool,
|
||||
},
|
||||
/// Help ▸ System Info: (label, value) rows.
|
||||
/// A face model file the user chose: what it is, its licence terms, and the "I accept" box.
|
||||
FaceModel {
|
||||
path: String,
|
||||
/// The engine's `faces.models.inspect` answer.
|
||||
info: serde_json::Value,
|
||||
accepted: bool,
|
||||
},
|
||||
SystemInfo {
|
||||
rows: Vec<(String, String)>,
|
||||
},
|
||||
|
||||
+2
-2
@@ -42,7 +42,7 @@ and prints the summary; `cargo xtask parity --write` refreshes the summary table
|
||||
| S. Export (EXP) | 15 | 1 | 2 | 0 | 7/7 (100%) | 7/7 (100%) |
|
||||
| T. Share (SHARE) | 0 | 0 | 0 | 4 | — | — |
|
||||
| U. Map & location (MAP) | 0 | 1 | 1 | 0 | — | 0/1 (0%) |
|
||||
| V. Preferences (PREF) | 6 | 0 | 3 | 2 | 1/1 (100%) | 4/4 (100%) |
|
||||
| V. Preferences (PREF) | 6 | 1 | 2 | 2 | 1/1 (100%) | 4/4 (100%) |
|
||||
| W. Cloud & AI infrastructure (CLOUD / AI) | 1 | 0 | 1 | 2 | — | — |
|
||||
| X. Cross-cutting behaviours (BEHAV) | 18 | 5 | 1 | 1 | 8/8 (100%) | 6/8 (75%) |
|
||||
| Y. Menus | 82 | 0 | 4 | 8 | 47/47 (100%) | 23/23 (100%) |
|
||||
@@ -482,7 +482,7 @@ under the UI command that wraps it (a no-op error); the UI command now wins.
|
||||
| LR-PREF-ACCOUNT | Account | OOS | 🚫 | | |
|
||||
| LR-PREF-INTERFACE | Interface options | P1 | ✅ | `cmd:app.settings` | filmstrip names/badges, grid badges (auto/always/never), square-grid names, navigator, info overlay |
|
||||
| LR-PREF-PERFORMANCE | GPU / performance | P1 | ✅ | `cmd:app.gpu`, `cmd:app.memoryBudget`, `cmd:app.settings` | GPU on/off, preview size (Automatic = the whole-frame loupe render at canvas size, up to 2560 px; or a 1600–5120 px size for it; zoomed views add a window at up to 100 %, issue #323), memory budget (a quarter of it is the idle per-view stage caches' share; the open photo's views are counted and reported, and cut only past four times that, the before side's and the zoom window's first — a Before/After at 1:1 on a 2800 px canvas holds about 1 GB of GPU stages, plus the one 288 MB copy of the original, which `ui.inspect` reports as `sharedSourceBytes`), thumbnail cache size in Settings |
|
||||
| LR-PREF-PEOPLE | Face recognition | P2 | ⬜ | | |
|
||||
| LR-PREF-PEOPLE | Face recognition | P2 | 🟡 | `cmd:faces.enable`, `cmd:faces.models.list`, `cmd:faces.models.inspect`, `cmd:faces.models.install`, `cmd:faces.models.remove`, `cmd:faces.models.select`, `cmd:dialog.faceModel`, `crates/faces/src/lib.rs`, `crates/ui-egui/src/panels/faces.rs` | Settings ▸ Faces: an off-by-default toggle and the face model list. Models are opt-in and never bundled (the YuNet detector, 233 KB, is the one exception, not wired yet): add one by picking or dropping a `.onnx`, which is recognised by hash or described from its shape, shows its licence terms and training-data note, and installs only after an "I accept" box. Nothing runs a model yet (no inference runtime), so nothing recognises faces yet |
|
||||
| LR-PREF-WATERMARK | Watermark settings | P1 | ✅ | `cmd:export.savePreset`, `crates/ui-egui/src/panels/dialogs.rs` | set in the Export dialog; kept with Export with Previous and in saved export presets |
|
||||
| LR-PREF-SHORTCUTS | Shortcut customisation | — | ✅ | `cmd:app.setShortcut`, `cmd:app.resetShortcuts`, `crates/ui-egui/src/panels/keymap.rs` | an extra (the reference app has no keymap editor): Help ▸ Keyboard Shortcuts (⌘/) lists every command; click a shortcut and press the new keys (Esc cancels), × removes, ↺ restores, Reset All; a key another command had moves over; saved in `ui.json`; native menu accelerators follow. Rating/label digits and secondary keys are listed read-only and yield to a key given to a command |
|
||||
| LR-PREF-TECHPREVIEW | Early-access toggles | P2 | ⬜ | | |
|
||||
|
||||
@@ -43,6 +43,7 @@ pub const TABLE: &[(&str, Class)] = &[
|
||||
("codecs", Class::Layer(1)),
|
||||
("meta", Class::Layer(1)),
|
||||
("develop", Class::Layer(1)),
|
||||
("faces", Class::Layer(1)),
|
||||
("scenes", Class::Layer(1)),
|
||||
("pipeline", Class::Layer(2)),
|
||||
("gpu", Class::Layer(3)),
|
||||
|
||||
Reference in new issue
Block a user