mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-30 04:57:22 -06:00
Merge remote-tracking branch 'origin/main' into fix-445
# Conflicts: # crates/engine/src/media.rs
This commit is contained in:
Brandon Thomas
@@ -65,19 +65,8 @@ impl CameraProfile {
|
||||
}
|
||||
}
|
||||
|
||||
/// LightCraft's configuration folder (settings, GPU marker, camera profiles).
|
||||
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"))
|
||||
}
|
||||
}
|
||||
/// LightCraft's configuration folder (settings, GPU marker, camera profiles); one helper for the app.
|
||||
pub use crate::config::config_dir;
|
||||
|
||||
/// Where camera profiles are read from and written to.
|
||||
pub fn dir() -> Option<PathBuf> {
|
||||
|
||||
@@ -0,0 +1,659 @@
|
||||
//! Denoise commands (`denoise.*`): the models (list, install, download, test, remove, choose), the settings, and the
|
||||
//! work (queue, cancel, status, pump, clear). The machinery they drive is [`crate::denoise`].
|
||||
//!
|
||||
//! Models live in the folder the host provides (`Session::denoise.models_dir`): one subfolder per model with
|
||||
//! `model.onnx`, `denoise-model.json` (its manifest) and `installed.json` (what the user accepted and the file's hash),
|
||||
//! plus `settings.json` beside them. Model files, archives 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::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, PoisonError, mpsc};
|
||||
|
||||
use lightcraft_catalog::PhotoId;
|
||||
use lightcraft_denoise::archive;
|
||||
use lightcraft_denoise::hash::sha256_file;
|
||||
use lightcraft_denoise::known::{self, Known};
|
||||
use lightcraft_denoise::manifest::{self, DenoiserManifest, MAX_MANIFEST_BYTES, MAX_MODEL_BYTES};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{CommandSpec, always, bad, bool_or, cmd, str_param};
|
||||
use crate::denoise::{Pace, RunOn, Settings, forget_gpu_set_ups, installed_models, read_capped, read_settings, write_atomic, write_settings};
|
||||
use crate::model_download::State;
|
||||
use crate::{EngineError, Result, Session};
|
||||
|
||||
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.denoise.models_dir.clone().ok_or_else(|| bad(cmd, "this build has nowhere to keep denoise models (the desktop app does)"))
|
||||
}
|
||||
|
||||
/// What a command says about a model.
|
||||
fn row(m: &DenoiserManifest, download_host: Option<String>, installed: bool, selected: bool, accepted: &Value) -> Value {
|
||||
json!({
|
||||
"id": m.id,
|
||||
"name": m.name,
|
||||
"version": m.version,
|
||||
"licence": m.licence,
|
||||
"provenance": m.provenance,
|
||||
"source": m.source,
|
||||
"sizeBytes": m.size_bytes,
|
||||
"sha256": m.sha256,
|
||||
"tile": m.tile,
|
||||
"known": known::find(&m.id).is_some(),
|
||||
// a pinned address LightCraft can fetch it from (the user presses Download), and which site that is
|
||||
"downloadHost": download_host,
|
||||
"installed": installed,
|
||||
"selected": selected,
|
||||
"accepted": accepted,
|
||||
})
|
||||
}
|
||||
|
||||
fn host_of(k: &Known) -> Option<String> {
|
||||
k.download.as_ref().map(|d| lightcraft_denoise::known::host(&d.url).to_string())
|
||||
}
|
||||
|
||||
fn list(s: &mut Session, _: &Value) -> Result<Value> {
|
||||
s.denoise_refresh_active(true);
|
||||
let dir = s.denoise.models_dir.clone();
|
||||
let st = dir.as_deref().map(read_settings).unwrap_or_default();
|
||||
let installed = dir.as_deref().map(installed_models).unwrap_or_default();
|
||||
let on_disk = |id: &str| installed.iter().find(|i| i.manifest.id == id);
|
||||
let chosen = |id: &str| st.model.as_deref() == Some(id);
|
||||
let mut models: Vec<Value> = Vec::new();
|
||||
for k in known::all() {
|
||||
let found = on_disk(&k.manifest.id);
|
||||
models.push(row(&k.manifest, host_of(&k), found.is_some(), chosen(&k.manifest.id), found.map_or(&Value::Null, |f| &f.accepted)));
|
||||
}
|
||||
for i in installed.iter().filter(|i| known::find(&i.manifest.id).is_none()) {
|
||||
models.push(row(&i.manifest, None, true, chosen(&i.manifest.id), &i.accepted));
|
||||
}
|
||||
Ok(json!({
|
||||
"dir": dir.as_ref().map(|d| d.display().to_string()),
|
||||
"productsDir": s.denoise_products_dir().map(|d| d.display().to_string()),
|
||||
"model": st.model,
|
||||
// whether this build can run denoise models (the `denoise` feature: our pure-Rust CPU/GPU runners)
|
||||
"runtime": cfg!(feature = "denoise"),
|
||||
"auto": st.auto(),
|
||||
"cacheGb": st.cache_bytes() >> 30,
|
||||
"threads": st.threads,
|
||||
"runOn": st.run_on().name(),
|
||||
"models": models,
|
||||
}))
|
||||
}
|
||||
|
||||
/// What a file the user pointed at is, as far as installing it goes.
|
||||
struct Source {
|
||||
manifest: DenoiserManifest,
|
||||
/// The `.onnx` inside the archive; `None` when the file is the `.onnx`.
|
||||
entry: Option<&'static str>,
|
||||
file_name: String,
|
||||
}
|
||||
|
||||
fn identify(path: &Path, c: &str, sha_known: Option<&str>) -> Result<Source> {
|
||||
let meta = std::fs::metadata(path).map_err(|e| bad(c, format!("cannot read `{}`: {e}", path.display())))?;
|
||||
if !meta.is_file() {
|
||||
return Err(bad(c, format!("`{}` is not a file", path.display())));
|
||||
}
|
||||
if meta.len() == 0 || meta.len() > MAX_MODEL_BYTES {
|
||||
return Err(bad(c, "a model file must be between 1 byte and 2 GiB"));
|
||||
}
|
||||
let sha256 = match sha_known {
|
||||
Some(h) => h.to_string(),
|
||||
None => sha256_file(path).map_err(|e| fail(&format!("could not read {}", path.display()), e))?,
|
||||
};
|
||||
let file_name = path.file_name().map(|n| n.to_string_lossy().into_owned()).unwrap_or_default();
|
||||
// an archive LightCraft knows (the file darktable publishes the model in)
|
||||
if let Some(k) = known::all().into_iter().find(|k| k.download.as_ref().is_some_and(|d| d.sha256 == sha256)) {
|
||||
return Ok(Source { manifest: k.manifest, entry: k.entry, file_name });
|
||||
}
|
||||
// the `.onnx` of a known model
|
||||
if let Some(m) = known::by_sha256(&sha256) {
|
||||
return Ok(Source { manifest: m, entry: None, file_name });
|
||||
}
|
||||
// an `.onnx` with its manifest beside it
|
||||
let sibling = path.with_file_name("denoise-model.json");
|
||||
if path.extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx"))
|
||||
&& let Some(bytes) = read_capped(&sibling, MAX_MANIFEST_BYTES)
|
||||
{
|
||||
let m = manifest::parse(&bytes).map_err(|e| bad(c, format!("{}: {e}", sibling.display())))?;
|
||||
return Ok(Source { manifest: m, entry: None, file_name });
|
||||
}
|
||||
Err(bad(
|
||||
c,
|
||||
"LightCraft does not know this file. A denoise model is an .onnx file with a denoise-model.json beside it that says how to use it \
|
||||
(see docs/denoise.md), or an archive LightCraft can fetch itself (`denoise.models.download`)",
|
||||
))
|
||||
}
|
||||
|
||||
/// Install the model in `src`. A model that passes its self-test becomes the one in use (`activate`, the default): installing
|
||||
/// a model is how someone says they want it. `verified` is for a download whose hash was already checked.
|
||||
fn install_file(
|
||||
dir: &Path,
|
||||
loader: &crate::denoise::Loader,
|
||||
run_on: RunOn,
|
||||
accepted_at: String,
|
||||
c: &str,
|
||||
src: &Path,
|
||||
acknowledged: bool,
|
||||
verified: bool,
|
||||
) -> Result<(Value, Arc<dyn crate::denoise::Model>)> {
|
||||
let ident = identify(src, c, None)?;
|
||||
let m = ident.manifest.clone();
|
||||
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 (see `denoise.models.list`) and pass `acknowledged: true` once they agree",
|
||||
));
|
||||
}
|
||||
let home = dir.join(&m.id);
|
||||
if home.join("model.onnx").exists() {
|
||||
return Err(bad(c, "that model is already installed; remove it before installing a replacement"));
|
||||
}
|
||||
std::fs::create_dir_all(&home).map_err(|e| fail("could not create the model's folder", e))?;
|
||||
let mut cleanup = InstallCleanup { home: home.clone(), armed: true };
|
||||
let (part, final_path) = (home.join("model.onnx.part"), home.join("model.onnx"));
|
||||
let _ = std::fs::remove_file(&part);
|
||||
let placed: Result<()> = match ident.entry {
|
||||
Some(entry) => {
|
||||
archive::extract(src, entry, &part, MAX_MODEL_BYTES).map(|_| ()).map_err(|e| fail("could not take the model out of the archive", e))
|
||||
}
|
||||
None if verified && std::fs::rename(src, &part).is_ok() => Ok(()),
|
||||
None => std::fs::copy(src, &part).map(|_| ()).map_err(|e| fail("could not copy the model (is the disk full?)", e)),
|
||||
};
|
||||
let sha256 = placed.and_then(|_| sha256_file(&part).map_err(|e| fail("could not check the model", e))).and_then(|h| match &m.sha256 {
|
||||
Some(want) if *want != h => Err(EngineError::Other("the model file is not the one its description names (its SHA-256 differs)".into())),
|
||||
_ => Ok(h),
|
||||
});
|
||||
let sha256 = match sha256 {
|
||||
Ok(h) => h,
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
// Test the staged file first. A failed replacement preserves the installed model.
|
||||
let model = loader(&part, &m).map_err(|e| bad(c, e))?;
|
||||
let test = match model.self_test(run_on) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(bad(c, e));
|
||||
}
|
||||
};
|
||||
// An installed id is immutable: installing again must never destroy a usable model.
|
||||
if final_path.exists() {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(bad(c, "that model is already installed; remove it before installing a replacement"));
|
||||
}
|
||||
std::fs::rename(&part, &final_path).map_err(|e| fail("could not finish the copy", e))?;
|
||||
write_atomic(&home.join("denoise-model.json"), &serde_json::to_vec_pretty(&m).map_err(|e| fail("manifest", e))?).map_err(EngineError::Other)?;
|
||||
let accepted = json!({
|
||||
"acceptedAt": accepted_at,
|
||||
"licence": m.licence.name,
|
||||
"commercial": m.licence.commercial,
|
||||
"fileName": ident.file_name,
|
||||
"sha256": sha256,
|
||||
"selfTest": test,
|
||||
});
|
||||
write_atomic(&home.join("installed.json"), &serde_json::to_vec_pretty(&accepted).map_err(|e| fail("record", e))?).map_err(EngineError::Other)?;
|
||||
cleanup.armed = false;
|
||||
Ok((json!({"installed": row(&m, known::find(&m.id).and_then(|k| host_of(&k)), true, false, &accepted), "model": m.id}), model))
|
||||
}
|
||||
|
||||
/// Removes only files created for this new installation, including on a worker panic.
|
||||
struct InstallCleanup {
|
||||
home: PathBuf,
|
||||
armed: bool,
|
||||
}
|
||||
impl Drop for InstallCleanup {
|
||||
fn drop(&mut self) {
|
||||
if self.armed {
|
||||
for name in ["model.onnx.part", "model.onnx", "denoise-model.json", "installed.json"] {
|
||||
let _ = std::fs::remove_file(self.home.join(name));
|
||||
}
|
||||
let _ = std::fs::remove_dir(&self.home);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct InstallJob {
|
||||
id: String,
|
||||
activate: bool,
|
||||
result: mpsc::Receiver<std::result::Result<(Value, Arc<dyn crate::denoise::Model>), String>>,
|
||||
}
|
||||
|
||||
fn activate_install(s: &mut Session, result: &Value, model: Arc<dyn crate::denoise::Model>, activate: bool) -> Result<Value> {
|
||||
let dir = models_dir(s, "denoise.models.install")?;
|
||||
let id = result["installed"]["id"].as_str().ok_or_else(|| bad("denoise.models.install", "missing installed id"))?;
|
||||
let mut st = read_settings(&dir);
|
||||
if activate {
|
||||
st.model = Some(id.to_string());
|
||||
write_settings(&dir, &st).map_err(EngineError::Other)?;
|
||||
}
|
||||
s.denoise.touch();
|
||||
s.denoise_refresh_active(true);
|
||||
if let Some(a) = &s.denoise.active
|
||||
&& a.id == id
|
||||
{
|
||||
*a.model.lock().unwrap_or_else(PoisonError::into_inner) = Some(model);
|
||||
}
|
||||
let mut v = result.clone();
|
||||
v["installed"]["selected"] = json!(st.model.as_deref() == Some(id));
|
||||
v["model"] = json!(st.model);
|
||||
Ok(v)
|
||||
}
|
||||
|
||||
fn start_install(s: &mut Session, src: PathBuf, id: String, activate: bool, verified: bool) -> Result<Value> {
|
||||
const C: &str = "denoise.models.install";
|
||||
if s.denoise.install.is_some() {
|
||||
return Err(bad(C, "another model is being installed"));
|
||||
}
|
||||
let dir = models_dir(s, C)?;
|
||||
let loader = s.denoise.loader.clone();
|
||||
let run_on = s.denoise.settings.run_on();
|
||||
let at = (s.clock)();
|
||||
let (tx, rx) = mpsc::channel();
|
||||
std::thread::Builder::new()
|
||||
.name("denoise-model-install".into())
|
||||
.stack_size(16 * 1024 * 1024)
|
||||
.spawn(move || {
|
||||
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
||||
let _held = crate::memory::work_gate().acquire(512 * 1024 * 1024);
|
||||
install_file(&dir, &loader, run_on, at, C, &src, true, verified).map_err(|e| plain(&e))
|
||||
}))
|
||||
.unwrap_or_else(|_| Err("the model installation stopped unexpectedly".into()));
|
||||
if verified {
|
||||
let _ = std::fs::remove_file(&src);
|
||||
}
|
||||
let _ = tx.send(outcome);
|
||||
})
|
||||
.map_err(|e| bad(C, format!("could not start installation: {e}")))?;
|
||||
s.denoise.install = Some(InstallJob { id: id.clone(), activate, result: rx });
|
||||
s.denoise.downloads.set_outcome(&id, State::Installing);
|
||||
Ok(json!({"started": id}))
|
||||
}
|
||||
|
||||
fn install(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.models.install";
|
||||
let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?;
|
||||
let acknowledged = bool_or(p, "acknowledged", false);
|
||||
let activate = bool_or(p, "activate", true);
|
||||
if !acknowledged {
|
||||
return Err(bad(C, "the licence has not been accepted: show its terms first"));
|
||||
}
|
||||
if bool_or(p, "background", false) {
|
||||
return start_install(s, PathBuf::from(path), "local-install".into(), activate, false);
|
||||
}
|
||||
if s.denoise.install.is_some() {
|
||||
return Err(bad(C, "another model is being installed"));
|
||||
}
|
||||
let (result, model) =
|
||||
install_file(&models_dir(s, C)?, &s.denoise.loader, s.denoise.settings.run_on(), (s.clock)(), C, Path::new(path), acknowledged, false)?;
|
||||
activate_install(s, &result, model, activate)
|
||||
}
|
||||
|
||||
/// Read only the small adjacent manifest; hashing and execution happen in the install worker.
|
||||
fn inspect(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.models.inspect";
|
||||
let path = str_param(p, "path").ok_or_else(|| bad(C, "missing `path`"))?;
|
||||
let src = Path::new(path);
|
||||
if !src.extension().is_some_and(|e| e.eq_ignore_ascii_case("onnx")) {
|
||||
return Err(bad(C, "choose an .onnx file with denoise-model.json beside it"));
|
||||
}
|
||||
let bytes = read_capped(&src.with_file_name("denoise-model.json"), MAX_MANIFEST_BYTES)
|
||||
.ok_or_else(|| bad(C, "denoise-model.json is missing or too large"))?;
|
||||
let m = manifest::parse(&bytes).map_err(|e| bad(C, e.to_string()))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
Ok(json!({"kind": "model", "domain": "denoise", "path": path, "fileName": src.file_name().map(|n| n.to_string_lossy()), "model": m,
|
||||
"sizeBytes": std::fs::metadata(src).ok().map(|m| m.len()), "alreadyInstalled": dir.join(&m.id).join("model.onnx").is_file()}))
|
||||
}
|
||||
|
||||
/// What an error says, without the command it came from (for a line shown to the user).
|
||||
fn plain(e: &EngineError) -> String {
|
||||
match e {
|
||||
EngineError::BadParams { msg, .. } => msg.clone(),
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Install every download that has arrived: each was accepted when it was started, so it is installed and chosen without
|
||||
/// another question. Called by `denoise.models.downloads` and by every frame's `denoise.pump`.
|
||||
pub(crate) fn finish_downloads(s: &mut Session) {
|
||||
if let Some(job) = &s.denoise.install {
|
||||
let outcome = match job.result.try_recv() {
|
||||
Ok(v) => Some(v),
|
||||
Err(mpsc::TryRecvError::Empty) => None,
|
||||
Err(mpsc::TryRecvError::Disconnected) => Some(Err("the installation worker stopped".into())),
|
||||
};
|
||||
if let Some(outcome) = outcome
|
||||
&& let Some(job) = s.denoise.install.take()
|
||||
{
|
||||
let state = match outcome {
|
||||
Ok((v, model)) => match activate_install(s, &v, model, job.activate) {
|
||||
Ok(_) => State::Installed,
|
||||
Err(e) => State::Failed(plain(&e)),
|
||||
},
|
||||
Err(e) => State::Failed(e),
|
||||
};
|
||||
s.denoise.downloads.set_outcome(&job.id, state);
|
||||
}
|
||||
}
|
||||
if s.denoise.install.is_none()
|
||||
&& let Some((id, path, _sha)) = s.denoise.downloads.finished().into_iter().next()
|
||||
&& let Err(e) = start_install(s, path, id.clone(), true, true)
|
||||
{
|
||||
s.denoise.downloads.set_outcome(&id, State::Failed(plain(&e)));
|
||||
}
|
||||
}
|
||||
|
||||
fn download(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.models.download";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
let k = known::find(id).ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?;
|
||||
let spec = k.download.clone().ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?;
|
||||
if !cfg!(feature = "denoise") {
|
||||
return Err(bad(C, "this build cannot run denoise models, so there is nothing to download for it"));
|
||||
}
|
||||
if !bool_or(p, "acknowledged", false) {
|
||||
return Err(bad(
|
||||
C,
|
||||
"the licence has not been accepted: show the user the model's terms (see `denoise.models.list`) and pass `acknowledged: true` once they agree",
|
||||
));
|
||||
}
|
||||
if installed_models(&dir).iter().any(|i| i.manifest.id == spec.id) {
|
||||
return Err(bad(C, "that model is already installed"));
|
||||
}
|
||||
let host = lightcraft_denoise::known::host(&spec.url).to_string();
|
||||
s.denoise.downloads.start(spec, &dir).map_err(|e| bad(C, e))?;
|
||||
Ok(json!({"started": id, "from": host}))
|
||||
}
|
||||
|
||||
fn download_row(id: &str, state: &State) -> Value {
|
||||
let from = known::find(id).and_then(|k| host_of(&k));
|
||||
match state {
|
||||
State::Running { bytes, total } => json!({"id": id, "state": "running", "bytes": bytes, "total": total, "from": from}),
|
||||
State::Done { path, .. } => json!({"id": id, "state": "done", "path": path.display().to_string(), "from": from}),
|
||||
State::Installing => json!({"id": id, "state": "installing", "from": from}),
|
||||
State::Installed => json!({"id": id, "state": "installed", "from": from}),
|
||||
State::Failed(why) => json!({"id": id, "state": "failed", "error": why, "from": from}),
|
||||
State::Cancelled => json!({"id": id, "state": "cancelled", "from": from}),
|
||||
}
|
||||
}
|
||||
|
||||
fn downloads(s: &mut Session, _: &Value) -> Result<Value> {
|
||||
finish_downloads(s);
|
||||
let all: Vec<Value> = s.denoise.downloads.snapshot().iter().map(|(id, st)| download_row(id, st)).collect();
|
||||
Ok(json!({"running": s.denoise.downloads.running(), "downloads": all}))
|
||||
}
|
||||
|
||||
fn download_cancel(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.models.downloadCancel";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
Ok(json!({"discarded": s.denoise.downloads.discard(id)}))
|
||||
}
|
||||
|
||||
fn test(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.models.test";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
let found = installed_models(&dir).into_iter().find(|i| i.manifest.id == id).ok_or_else(|| bad(C, "that model is not installed"))?;
|
||||
let run_on = s.denoise.settings.run_on();
|
||||
let result = (s.denoise.loader)(&found.onnx, &found.manifest).and_then(|m| m.self_test(run_on));
|
||||
let (ok, detail) = match &result {
|
||||
Ok(v) => (true, v.clone()),
|
||||
Err(e) => (false, json!(e)),
|
||||
};
|
||||
// keep the latest outcome with the model's record
|
||||
let mut record = found.accepted;
|
||||
if let Some(o) = record.as_object_mut() {
|
||||
o.insert("selfTest".into(), if ok { detail.clone() } else { json!({"ok": false, "error": detail}) });
|
||||
write_atomic(&found.onnx.with_file_name("installed.json"), &serde_json::to_vec_pretty(&record).map_err(|e| fail("record", e))?)
|
||||
.map_err(EngineError::Other)?;
|
||||
}
|
||||
Ok(json!({"id": id, "ok": ok, "result": detail}))
|
||||
}
|
||||
|
||||
fn remove(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.models.remove";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
if !lightcraft_denoise::licence::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("denoise-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.model.as_deref() == Some(id) {
|
||||
st.model = installed_models(&dir).into_iter().map(|i| i.manifest.id).find(|m| m != id);
|
||||
write_settings(&dir, &st).map_err(EngineError::Other)?;
|
||||
}
|
||||
s.denoise.touch();
|
||||
s.denoise_refresh_active(true);
|
||||
Ok(json!({"removed": id, "model": st.model}))
|
||||
}
|
||||
|
||||
fn select(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.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)
|
||||
{
|
||||
return Err(bad(C, "that denoise model is not installed"));
|
||||
}
|
||||
let mut st = read_settings(&dir);
|
||||
st.model = id;
|
||||
write_settings(&dir, &st).map_err(EngineError::Other)?;
|
||||
s.denoise.touch();
|
||||
s.denoise_refresh_active(true);
|
||||
Ok(json!({"model": st.model}))
|
||||
}
|
||||
|
||||
fn settings(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.settings";
|
||||
let dir = models_dir(s, C)?;
|
||||
let mut st: Settings = read_settings(&dir);
|
||||
if let Some(v) = p.get("auto") {
|
||||
st.auto = Some(v.as_bool().ok_or_else(|| bad(C, "`auto` must be true or false"))?);
|
||||
}
|
||||
if let Some(v) = p.get("cacheGb") {
|
||||
let n =
|
||||
v.as_u64().filter(|n| (1..=10_000).contains(n)).ok_or_else(|| bad(C, "`cacheGb` must be a whole number of gigabytes from 1 to 10000"))?;
|
||||
st.cache_gb = Some(n as u32);
|
||||
}
|
||||
if let Some(v) = p.get("threads") {
|
||||
st.threads = match v {
|
||||
Value::Null => None,
|
||||
v => Some(v.as_u64().filter(|n| (1..=64).contains(n)).ok_or_else(|| bad(C, "`threads` must be from 1 to 64, or null"))? as u32),
|
||||
};
|
||||
}
|
||||
let run_on = match p.get("runOn") {
|
||||
None => None,
|
||||
Some(v) => Some(v.as_str().and_then(RunOn::parse).ok_or_else(|| bad(C, "`runOn` must be auto, gpu or cpu"))?),
|
||||
};
|
||||
if let Some(r) = run_on {
|
||||
st.run_on = Some(r.name().to_string());
|
||||
st.gpu = None;
|
||||
}
|
||||
write_settings(&dir, &st).map_err(EngineError::Other)?;
|
||||
if run_on.is_some() {
|
||||
// a choice of where denoise runs is also a go-ahead to try the card again: a set-up that crashed or gave up
|
||||
// before is forgotten, and the loaded model sets the card up afresh for the next photo
|
||||
forget_gpu_set_ups(&dir);
|
||||
if let Some(a) = &s.denoise.active {
|
||||
*a.model.lock().unwrap_or_else(PoisonError::into_inner) = None;
|
||||
}
|
||||
}
|
||||
s.denoise.touch();
|
||||
s.denoise_refresh_active(true);
|
||||
// a smaller limit takes effect now
|
||||
s.denoise_evict(None);
|
||||
Ok(json!({"auto": st.auto(), "cacheGb": st.cache_bytes() >> 30, "threads": st.threads, "runOn": st.run_on().name()}))
|
||||
}
|
||||
|
||||
fn ids_param(p: &Value) -> Option<Vec<PhotoId>> {
|
||||
p.get("ids").and_then(Value::as_array).map(|a| a.iter().filter_map(Value::as_u64).map(PhotoId).collect())
|
||||
}
|
||||
|
||||
fn status_json(s: &Session, photo: Option<PhotoId>) -> Value {
|
||||
let (files, bytes) = s.denoise_cache_usage();
|
||||
let st = &s.denoise.settings;
|
||||
let model = s.denoise.active.as_ref().map(|a| json!({"id": a.id, "name": a.manifest.name}));
|
||||
let mut v = json!({
|
||||
"enabled": s.denoise.active.is_some(),
|
||||
"model": model,
|
||||
"runtime": cfg!(feature = "denoise"),
|
||||
"productsDir": s.denoise_products_dir().map(|d| d.display().to_string()),
|
||||
"queued": s.denoise.queued(),
|
||||
"running": s.denoise_running_json(),
|
||||
"ready": s.media.denoise.len(),
|
||||
"made": s.denoise_made(),
|
||||
"auto": st.auto(),
|
||||
"runOn": st.run_on().name(),
|
||||
"device": s.denoise_device(),
|
||||
"cache": {"files": files, "bytes": bytes, "limitBytes": st.cache_bytes()},
|
||||
"generation": s.denoise.generation,
|
||||
});
|
||||
if let (Some(id), Some(o)) = (photo, v.as_object_mut()) {
|
||||
o.insert("photo".into(), serde_json::to_value(s.denoise_photo_state(id)).unwrap_or(Value::Null));
|
||||
}
|
||||
v
|
||||
}
|
||||
|
||||
/// A per-photo switch, independent of the selection and Auto Sync. Turning on starts at 50%.
|
||||
fn toggle(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.toggle";
|
||||
let id = match p.get("id") {
|
||||
Some(v) => PhotoId(v.as_u64().ok_or_else(|| bad(C, "`id` must be a photo id"))?),
|
||||
None => s.active().ok_or_else(|| bad(C, "select a raw photo"))?,
|
||||
};
|
||||
let photo = s.catalog.photo(id).ok_or_else(|| bad(C, "no such photo"))?;
|
||||
if !crate::denoise::eligible(photo) {
|
||||
return Err(bad(C, "AI Denoise applies to raw photos developed from their sensor data"));
|
||||
}
|
||||
let mut d = (*s.develop_of(id).ok_or_else(|| bad(C, "no such photo"))?).clone();
|
||||
let on = match p.get("enabled") {
|
||||
Some(v) => v.as_bool().ok_or_else(|| bad(C, "`enabled` must be true or false"))?,
|
||||
None => !d.enhance.denoise_enabled(),
|
||||
};
|
||||
if on {
|
||||
if !d.enhance.denoise.is_finite() {
|
||||
return Err(bad(C, "the photo has an invalid Denoise amount"));
|
||||
}
|
||||
s.denoise_refresh_active(true);
|
||||
if s.denoise.active.is_none() {
|
||||
return Err(bad(C, "AI Denoise needs a model: install one in Settings > AI Denoise"));
|
||||
}
|
||||
}
|
||||
// The switch is its own setting: the Amount stays where it is (0 included) when it is flipped, so flipping compares
|
||||
// with and without. A photo that never had an Amount starts at 50 the first time.
|
||||
let amount = if on && d.enhance.denoise_on.is_none() && d.enhance.denoise <= 0.0 { 50.0 } else { d.enhance.denoise };
|
||||
if d.enhance.denoise != amount || d.enhance.denoise_on != Some(on) {
|
||||
s.end_interaction()?;
|
||||
d.enhance.denoise = amount;
|
||||
d.enhance.denoise_on = Some(on);
|
||||
let op = s.develop_op(id, d, "AI Denoise").ok_or_else(|| bad(C, "no such photo"))?;
|
||||
s.commit("AI Denoise", op)?;
|
||||
}
|
||||
if on {
|
||||
s.denoise_enqueue(&[id], true, false);
|
||||
} else {
|
||||
s.denoise_cancel(Some(id), false);
|
||||
}
|
||||
Ok(json!({"id": id.0, "enabled": on, "amount": amount}))
|
||||
}
|
||||
|
||||
fn queue(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "denoise.queue";
|
||||
s.denoise_refresh_active(true);
|
||||
if s.denoise.active.is_none() {
|
||||
return Err(bad(C, "no denoise model is chosen: install one first (`denoise.models.download` or `denoise.models.install`)"));
|
||||
}
|
||||
if s.denoise_products_dir().is_none() {
|
||||
return Err(bad(C, "denoise keeps its pictures in the library's folder: open a library on disk first"));
|
||||
}
|
||||
let ids = match ids_param(p) {
|
||||
Some(ids) => ids,
|
||||
None => match str_param(p, "scope").unwrap_or("selected") {
|
||||
"selected" => s.selection.active.iter().chain(s.selection.ids.iter()).copied().collect(),
|
||||
"visible" => s.visible_cloned().to_vec(),
|
||||
"withAmount" => s.catalog.photos().filter(|ph| ph.develop.denoise_amount() > 0.0).map(|ph| ph.id).collect(),
|
||||
other => return Err(bad(C, format!("unknown scope `{other}`: use selected, visible or withAmount"))),
|
||||
},
|
||||
};
|
||||
let queued = s.denoise_enqueue(&ids, false, bool_or(p, "retry", false));
|
||||
let mut v = status_json(s, None);
|
||||
if let Some(o) = v.as_object_mut() {
|
||||
o.insert("added".into(), json!(queued));
|
||||
}
|
||||
Ok(v)
|
||||
}
|
||||
|
||||
fn cancel(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
let photo = p.get("id").and_then(Value::as_u64).map(PhotoId);
|
||||
let n = s.denoise_cancel(photo, bool_or(p, "all", photo.is_none()));
|
||||
Ok(json!({"cancelled": n}))
|
||||
}
|
||||
|
||||
fn status(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
s.denoise_refresh_active(false);
|
||||
Ok(status_json(s, p.get("id").and_then(Value::as_u64).map(PhotoId)))
|
||||
}
|
||||
|
||||
/// `denoise.pump {pace?}`: cheap to call every frame. Takes in what the background job has finished and keeps one going.
|
||||
fn pump(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
s.denoise_pump(Pace::parse(p.get("pace").and_then(Value::as_str)));
|
||||
Ok(json!({
|
||||
"active": s.denoise.active.is_some(),
|
||||
"running": s.denoise_running_json(),
|
||||
"queued": s.denoise.queued(),
|
||||
"ready": s.media.denoise.len(),
|
||||
"generation": s.denoise.generation,
|
||||
"downloading": s.denoise.downloads.running(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn clear(s: &mut Session, _: &Value) -> Result<Value> {
|
||||
let (files, bytes) = s.denoise_clear();
|
||||
Ok(json!({"deleted": files, "bytes": bytes}))
|
||||
}
|
||||
|
||||
pub fn specs() -> Vec<CommandSpec> {
|
||||
vec![
|
||||
cmd!(query "denoise.models.list", "Denoise Models", [], None, "{} → {dir, productsDir, model, runtime, auto, cacheGb, threads, runOn, models: [{id, name, version, licence{name, commercial, url, notice}, provenance, source, sizeBytes, sha256, tile, known, downloadHost, installed, selected, accepted}]}", always, list),
|
||||
cmd!(query "denoise.models.inspect", "Inspect Denoise Model", [], None, "{path} → model terms from an adjacent manifest; never executes the file", always, inspect),
|
||||
cmd!(query "denoise.models.install", "Install Denoise Model", [], None, "{path, acknowledged: true, activate?: true, background?: false} → {installed, model} or {started} — install a denoise model from a file: an .onnx with a denoise-model.json beside it, or an archive LightCraft knows (the darktable `.dtmodel`). `acknowledged` must be true: the user has been shown the model's licence and accepted it. Unless `activate` is false it becomes the model in use", always, install),
|
||||
cmd!(query "denoise.models.download", "Download Denoise Model", [], None, "{id, acknowledged: true} → {started, from} — fetch a model LightCraft has a pinned address for (see `downloadHost` in the list) in the background over HTTPS with lightcraft-fetch. `acknowledged` must be true: the user has been shown the model's terms and accepted them. It is checked against its size and SHA-256, installed and chosen by itself; `denoise.models.downloads` shows how far it is", always, download),
|
||||
cmd!(query "denoise.models.downloads", "Denoise Model Downloads", [], None, "{} → {running, downloads: [{id, state: running | done | installed | failed | cancelled, bytes, total, error, from}]} — also installs any download that has arrived", always, downloads),
|
||||
cmd!(query "denoise.models.downloadCancel", "Cancel Denoise Model Download", [], None, "{id} → {discarded} — stop a download, or delete a finished one that was not installed", always, download_cancel),
|
||||
cmd!(query "denoise.models.test", "Test Denoise Model", [], None, "{id} → {ok, result} — load an installed model and check it gives sensible pictures; needs the denoise runtime", always, test),
|
||||
cmd!(query "denoise.models.remove", "Remove Denoise Model", [], None, "{id} → {removed, model} — delete an installed model (its cached pictures stay until they are evicted)", always, remove),
|
||||
cmd!(query "denoise.models.select", "Choose Denoise Model", [], None, "{id: installed model | null} → {model} — null switches denoise off", always, select),
|
||||
cmd!(query "denoise.settings", "Denoise Settings", [], None, "{auto?: bool, cacheGb?: 1..10000, threads?: 1..64 | null, runOn?: auto | gpu | cpu} → the settings — `auto`: make the picture of the photos being looked at that have a Denoise amount; `cacheGb`: how much the cached pictures may take; `threads`: most tiles run at once; `runOn`: where the model runs — `auto` (the default) the graphics card where it is faster than the processor (both timed on this computer when the card is set up), `gpu` the card whenever it can run the model, `cpu` the processor. Setting `runOn` also lets the card be tried again after a set-up that failed or closed LightCraft", always, settings),
|
||||
cmd!(
|
||||
"denoise.toggle",
|
||||
"AI Denoise",
|
||||
[],
|
||||
None,
|
||||
"{id?: photoId, enabled?: bool} → {id, enabled, amount} — turn AI Denoise on (50% initially) or off for one raw photo, independent of selection and Auto Sync; needs an installed model to turn on",
|
||||
always,
|
||||
toggle
|
||||
),
|
||||
cmd!(query "denoise.queue", "Denoise Photos", [], None, "{ids?: [photoId], scope?: selected | visible | withAmount, retry?: false} → {added, …status} — make the denoised pictures of raw photos that have none, in the background. `retry` tries photos that failed before", always, queue),
|
||||
cmd!(query "denoise.cancel", "Cancel Denoise", [], None, "{id?, all?} → {cancelled} — stop the photo being made and/or forget the ones waiting (all of them when `id` is omitted)", always, cancel),
|
||||
cmd!(query "denoise.status", "Denoise Status", [], None, "{id?} → {enabled, model, runtime, productsDir, queued, running: {photo, done, total, seconds} | null, ready, made, auto, runOn, device: {kind: gpu | cpu | pending | none, adapter?, reason?, cardMs?, cpuMs?, fellBack?}, cache: {files, bytes, limitBytes}, generation, photo?: {state: notApplicable | noModel | ready | queued | running | failed | idle, …}}", always, status),
|
||||
cmd!(query "denoise.pump", "Denoise Pump", [], None, "{pace?: pause | light | normal | full} → {active, running, queued, ready, generation, downloading} — called every frame by the app: takes in finished work and keeps one photo going", always, pump),
|
||||
cmd!(query "denoise.clear", "Clear Denoise Cache", [], None, "{} → {deleted, bytes} — delete every cached denoised picture of this library (they are made again when wanted)", always, clear),
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
//! `faces.detect`: find faces with the YuNet detector the user has downloaded (Settings > Faces).
|
||||
//!
|
||||
//! The photo is rendered upright and uncropped with default settings (so edits, crops and rotations you have
|
||||
//! made do not matter), the detector runs on it, and the faces come back as fractions of the photo's width
|
||||
//! and height, the same frame face regions from XMP use. Unless `apply` is false, they are added to the photo as
|
||||
//! unnamed face regions in one undoable step; earlier detections are replaced, and regions that came from XMP (or
|
||||
//! that you drew or named) are never touched. Like every region edit it is catalog-only: no sidecar is written.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex, PoisonError};
|
||||
use std::time::Instant;
|
||||
|
||||
use lightcraft_catalog::Op;
|
||||
use lightcraft_develop::DevelopSettings;
|
||||
use lightcraft_faces::known::YUNET_ID;
|
||||
use lightcraft_faces::yunet::{Detector, Face, Options};
|
||||
use lightcraft_meta::{Region, RegionKind};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{CommandSpec, bad, cmd, has_selection};
|
||||
use crate::{Result, Session};
|
||||
|
||||
const C: &str = "faces.detect";
|
||||
/// Regions made by this command say so in their description; only those are replaced on a new run.
|
||||
pub(crate) const MARK: &str = "Detected by ";
|
||||
pub(crate) const LABEL: &str = "YuNet 2023mar";
|
||||
/// Long edge of the picture the detector looks at (it shrinks it to its own 640 anyway).
|
||||
const LOOK_EDGE: usize = 1280;
|
||||
/// YuNet is 232 KB; a file this much larger in its folder is not it.
|
||||
const MAX_MODEL_BYTES: u64 = 8 << 20;
|
||||
/// How much a detection must overlap a region the photo already has to count as the same face.
|
||||
const EXISTING_IOU: f64 = 0.3;
|
||||
|
||||
/// What to tell someone who has not downloaded the detector yet.
|
||||
const NOT_INSTALLED: &str = "finding faces needs the YuNet face detector (232 KB, MIT): download it in Settings > Faces";
|
||||
|
||||
/// The installed YuNet, loaded once and kept (a new one is loaded when the model folder changes).
|
||||
/// (also used by the background scan, which aligns faces by the detector's landmarks and finds faces in photos that have none)
|
||||
pub(crate) fn detector(s: &Session) -> std::result::Result<Arc<Detector>, String> {
|
||||
static LOADED: Mutex<Option<(PathBuf, Arc<Detector>)>> = Mutex::new(None);
|
||||
let home = s.face_models_dir.as_deref().ok_or("this build has nowhere to keep face models (the desktop app does)")?.join(YUNET_ID);
|
||||
let path = home.join("model.onnx");
|
||||
// installing writes the manifest last, so a folder without one is not (yet) a model
|
||||
if !path.is_file() || !home.join("face-model.json").is_file() {
|
||||
return Err(NOT_INSTALLED.into());
|
||||
}
|
||||
let mut loaded = LOADED.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some((p, d)) = loaded.as_ref()
|
||||
&& *p == path
|
||||
{
|
||||
return Ok(d.clone());
|
||||
}
|
||||
let len = std::fs::metadata(&path).map_err(|e| format!("could not read the detector: {e}"))?.len();
|
||||
if len > MAX_MODEL_BYTES {
|
||||
return Err("the detector's file is far bigger than YuNet: remove it in Settings > Faces and download it again".into());
|
||||
}
|
||||
let bytes = std::fs::read(&path).map_err(|e| format!("could not read the detector: {e}"))?;
|
||||
let d = Arc::new(Detector::new(&bytes).map_err(|e| format!("{e}: remove the detector in Settings > Faces and download it again"))?);
|
||||
*loaded = Some((path, d.clone()));
|
||||
Ok(d)
|
||||
}
|
||||
|
||||
fn face_json(f: &Face) -> Value {
|
||||
json!({
|
||||
"rect": {"x0": f.x0, "y0": f.y0, "x1": f.x1, "y1": f.y1},
|
||||
"score": f.score,
|
||||
"landmarks": f.landmarks.iter().map(|(x, y)| json!([x, y])).collect::<Vec<_>>(),
|
||||
})
|
||||
}
|
||||
|
||||
/// A detected face that lies on a region the photo already has (from XMP, drawn or named) is that region, not a new face.
|
||||
fn overlaps_existing(f: &Face, regions: &[Region]) -> bool {
|
||||
regions.iter().any(|r| {
|
||||
let (w, h) = (
|
||||
(r.rect.x1.min(f64::from(f.x1)) - r.rect.x0.max(f64::from(f.x0))).max(0.0),
|
||||
(r.rect.y1.min(f64::from(f.y1)) - r.rect.y0.max(f64::from(f.y0))).max(0.0),
|
||||
);
|
||||
let inter = w * h;
|
||||
let union = (r.rect.x1 - r.rect.x0) * (r.rect.y1 - r.rect.y0) + f64::from((f.x1 - f.x0) * (f.y1 - f.y0)) - inter;
|
||||
union > 0.0 && inter / union >= EXISTING_IOU
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn is_detected(r: &Region) -> bool {
|
||||
r.description.as_deref().is_some_and(|d| d.starts_with(MARK))
|
||||
}
|
||||
|
||||
/// The photo as the faces see it: upright, uncropped, default settings, as 8-bit RGB with its size, its long edge at
|
||||
/// most `edge`.
|
||||
pub(crate) fn render_rgb(s: &mut Session, id: lightcraft_catalog::PhotoId, edge: usize) -> std::result::Result<(Vec<u8>, usize, usize), String> {
|
||||
let settings = DevelopSettings::default();
|
||||
let job = s.preview_job(id, edge, edge, false, &settings).ok_or("no such photo")?;
|
||||
let image = job.run().rendered?.image;
|
||||
let rgb: Vec<u8> = image.data.iter().flat_map(|px| [px[0], px[1], px[2]]).collect();
|
||||
Ok((rgb, image.width, image.height))
|
||||
}
|
||||
|
||||
fn detect(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
let apply = p.get("apply").and_then(Value::as_bool).unwrap_or(true);
|
||||
let defaults = Options::default();
|
||||
let opts = Options {
|
||||
score: p.get("score").and_then(Value::as_f64).map_or(defaults.score, |v| v as f32),
|
||||
nms_iou: p.get("nmsIou").and_then(Value::as_f64).map_or(defaults.nms_iou, |v| v as f32),
|
||||
max_faces: p.get("maxFaces").and_then(Value::as_u64).map_or(defaults.max_faces, |v| usize::try_from(v).unwrap_or(defaults.max_faces)),
|
||||
};
|
||||
let detector = detector(s).map_err(|e| bad(C, e))?;
|
||||
let started = Instant::now();
|
||||
let (mut results, mut ops) = (Vec::new(), Vec::new());
|
||||
for id in s.targets(p) {
|
||||
let (rgb, width, height) = match render_rgb(s, id, LOOK_EDGE) {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
results.push(json!({"id": id.0, "error": e}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let faces = match detector.detect(&rgb, width, height, &opts) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
results.push(json!({"id": id.0, "error": e.to_string()}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if apply && let Some(photo) = s.catalog.photo(id) {
|
||||
let mut meta = photo.meta.clone();
|
||||
meta.regions.retain(|r| !is_detected(r));
|
||||
let fresh: Vec<&Face> = faces.iter().filter(|f| !overlaps_existing(f, &meta.regions)).collect();
|
||||
for f in fresh {
|
||||
meta.regions.push(Region {
|
||||
rect: lightcraft_geom::Rect { x0: f64::from(f.x0), y0: f64::from(f.y0), x1: f64::from(f.x1), y1: f64::from(f.y1) },
|
||||
kind: RegionKind::Face,
|
||||
name: None,
|
||||
description: Some(format!("{MARK}{LABEL}")),
|
||||
});
|
||||
}
|
||||
ops.push(Op::SetMeta { id, meta: Box::new(meta) });
|
||||
}
|
||||
results.push(json!({"id": id.0, "faces": faces.iter().map(face_json).collect::<Vec<_>>()}));
|
||||
}
|
||||
let applied = !ops.is_empty();
|
||||
if applied {
|
||||
s.commit("Detect Faces", Op::Batch { ops })?;
|
||||
// regions are not written to XMP, so a sidecar must not be rewritten for this
|
||||
s.skip_auto_write = true;
|
||||
}
|
||||
Ok(json!({"detector": LABEL, "applied": applied, "photos": results, "ms": started.elapsed().as_secs_f64() * 1000.0}))
|
||||
}
|
||||
|
||||
pub fn specs() -> Vec<CommandSpec> {
|
||||
vec![cmd!(
|
||||
"faces.detect",
|
||||
"Detect Faces",
|
||||
["Photo"],
|
||||
None,
|
||||
"{ids?, apply?: bool (default true), score?: 0.6, nmsIou?: 0.3, maxFaces?} → {photos: [{id, faces: [{rect, score, landmarks}]}]}. Finds faces with the YuNet detector (download it first: `faces.models.download {id: \"yunet-2023mar\"}`) and, unless `apply` is false, adds them as unnamed face regions (one undo step), replacing earlier detections but never regions from XMP or ones you named; the sidecar is not touched",
|
||||
has_selection,
|
||||
detect
|
||||
)]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn face(x0: f32, y0: f32, x1: f32, y1: f32) -> Face {
|
||||
Face { x0, y0, x1, y1, score: 0.9, landmarks: [(0.0, 0.0); 5] }
|
||||
}
|
||||
fn region(x0: f64, y0: f64, x1: f64, y1: f64) -> Region {
|
||||
Region { rect: lightcraft_geom::Rect { x0, y0, x1, y1 }, kind: RegionKind::Face, name: Some("Ann".into()), description: None }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_detection_on_an_existing_region_is_not_a_new_face() {
|
||||
let existing = [region(0.40, 0.20, 0.60, 0.50), region(0.10, 0.10, 0.20, 0.20)];
|
||||
// the same face, a bit larger or shifted (a Lightroom box is not YuNet's box)
|
||||
assert!(overlaps_existing(&face(0.38, 0.18, 0.64, 0.54), &existing));
|
||||
assert!(overlaps_existing(&face(0.42, 0.22, 0.58, 0.48), &existing));
|
||||
// another face, a sliver of overlap, and no regions at all
|
||||
assert!(!overlaps_existing(&face(0.70, 0.20, 0.90, 0.50), &existing));
|
||||
assert!(!overlaps_existing(&face(0.59, 0.49, 0.80, 0.80), &existing));
|
||||
assert!(!overlaps_existing(&face(0.0, 0.0, 0.0, 0.0), &existing), "an empty box overlaps nothing");
|
||||
assert!(!overlaps_existing(&face(0.4, 0.2, 0.6, 0.5), &[]));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,511 @@
|
||||
//! 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::catalog::{self, Catalog};
|
||||
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_with};
|
||||
use lightcraft_faces::{known, onnx};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{CommandSpec, always, bad, cmd, str_param};
|
||||
use crate::face_download::{self, State};
|
||||
use crate::{EngineError, Result, Session};
|
||||
|
||||
#[derive(Default, Serialize, Deserialize)]
|
||||
#[serde(default, rename_all = "camelCase")]
|
||||
pub(super) struct FaceSettings {
|
||||
pub(super) enabled: bool,
|
||||
/// The installed recogniser in use.
|
||||
pub(super) 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)
|
||||
}
|
||||
|
||||
pub(super) 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))?)
|
||||
}
|
||||
|
||||
pub(super) struct Installed {
|
||||
pub(super) manifest: ModelManifest,
|
||||
pub(super) 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.
|
||||
pub(super) 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
|
||||
}
|
||||
|
||||
/// Read the user's catalog (`catalog.json` in the models folder) anew: a few kilobytes, so it is always current.
|
||||
fn refresh_catalog(s: &mut Session) {
|
||||
let Some(dir) = s.face_models_dir.clone() else { return };
|
||||
let file = dir.join("catalog.json");
|
||||
s.face_catalog = match read_capped(&file, catalog::MAX_CATALOG_BYTES) {
|
||||
Some(bytes) => catalog::parse(&bytes),
|
||||
None if file.is_file() => catalog::parse(&vec![b' '; catalog::MAX_CATALOG_BYTES + 1]),
|
||||
None => Catalog::default(),
|
||||
};
|
||||
}
|
||||
|
||||
/// The manifests of the catalog's models (what a file's hash is matched against, besides the built-in ones).
|
||||
fn catalog_manifests(s: &Session) -> Vec<ModelManifest> {
|
||||
s.face_catalog.entries.iter().map(|e| e.manifest.clone()).collect()
|
||||
}
|
||||
|
||||
/// How to fetch a model: the built-in address, or the one in the user's catalog.
|
||||
fn download_spec(s: &Session, id: &str) -> Option<known::Download> {
|
||||
known::download(id).or_else(|| s.face_catalog.entries.iter().find(|e| e.manifest.id == id).and_then(catalog::Entry::download))
|
||||
}
|
||||
|
||||
fn row(s: &Session, m: &ModelManifest, installed: bool, selected: bool, accepted: &Value) -> Value {
|
||||
let from_catalog = s.face_catalog.entries.iter().any(|e| e.manifest.id == m.id);
|
||||
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,
|
||||
// how fast it is as a multiple of a ResNet-100 model (computers differ, ratios between models hardly do), and
|
||||
// that in words; null when nobody has measured it
|
||||
"speed": m.speed,
|
||||
"speedText": m.speed.map(manifest::describe_speed).filter(|t| !t.is_empty()),
|
||||
"known": from_catalog || known::all().iter().any(|k| k.id == m.id),
|
||||
// from the user's own catalog file
|
||||
"fromCatalog": from_catalog,
|
||||
// a pinned address LightCraft can fetch it from (the user presses Download), and which site that is
|
||||
"downloadHost": download_spec(s, &m.id).map(|d| known::host(&d.url).to_string()),
|
||||
"installed": installed,
|
||||
"selected": selected,
|
||||
"accepted": accepted,
|
||||
})
|
||||
}
|
||||
|
||||
struct Inspection {
|
||||
sha256: String,
|
||||
size: u64,
|
||||
file_name: String,
|
||||
suggestion: Suggestion,
|
||||
}
|
||||
|
||||
/// Look at a model file. `known_sha` is its SHA-256 when something already checked it (a finished download), so a
|
||||
/// file of hundreds of megabytes is not read again just to be hashed.
|
||||
fn inspect_file(path: &str, cmd: &str, known_sha: Option<&str>, extra: &[ModelManifest]) -> 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 512 MiB"));
|
||||
}
|
||||
let sha256 = match known_sha {
|
||||
Some(h) => h.to_string(),
|
||||
None => 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_with(&info, &sha256, meta.len(), &file_name, extra),
|
||||
Err(e) => Suggestion::Unsupported(e.to_string()),
|
||||
};
|
||||
Ok(Inspection { sha256, size: meta.len(), file_name, suggestion })
|
||||
}
|
||||
|
||||
fn list(s: &mut Session, _: &Value) -> Result<Value> {
|
||||
refresh_catalog(s);
|
||||
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(s, &k, found.is_some(), settings.embedder.as_deref() == Some(k.id.as_str()), found.map_or(&Value::Null, |f| &f.accepted)));
|
||||
}
|
||||
// the user's own catalog
|
||||
for e in &s.face_catalog.entries {
|
||||
let found = on_disk(&e.manifest.id);
|
||||
let chosen = settings.embedder.as_deref() == Some(e.manifest.id.as_str());
|
||||
models.push(row(s, &e.manifest, found.is_some(), chosen, found.map_or(&Value::Null, |f| &f.accepted)));
|
||||
}
|
||||
for i in &installed {
|
||||
let listed = known::all().iter().any(|k| k.id == i.manifest.id) || s.face_catalog.entries.iter().any(|e| e.manifest.id == i.manifest.id);
|
||||
if !listed {
|
||||
models.push(row(s, &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,
|
||||
// Native builds always have the pure-Rust runner; model files/background scanning remain native-only.
|
||||
"runtime": cfg!(not(target_arch = "wasm32")),
|
||||
"models": models,
|
||||
// the user's own list of models to download: where it goes, how many it holds, and what was wrong with the rest
|
||||
"catalog": {
|
||||
"path": dir.as_ref().map(|d| d.join("catalog.json").display().to_string()),
|
||||
"models": s.face_catalog.entries.len(),
|
||||
"errors": s.face_catalog.errors,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
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`"))?;
|
||||
refresh_catalog(s);
|
||||
let ins = inspect_file(path, C, None, &catalog_manifests(s))?;
|
||||
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(s, 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 activate = p.get("activate").and_then(Value::as_bool).unwrap_or(true);
|
||||
refresh_catalog(s);
|
||||
let ins = inspect_file(path, C, None, &catalog_manifests(s))?;
|
||||
install_file(s, C, Path::new(path), ins, acknowledged, activate, false)
|
||||
}
|
||||
|
||||
/// Install the inspected model file `src`. A recognition model that passes its self-test becomes the one in use and
|
||||
/// recognition is switched on (`activate`, the default): installing a model is how someone says they want it.
|
||||
/// `verified` is for a download that was already checked against its recorded hash: it is moved into place instead
|
||||
/// of being copied and read again.
|
||||
fn install_file(s: &mut Session, c: &str, src: &Path, ins: Inspection, acknowledged: bool, activate: bool, verified: bool) -> Result<Value> {
|
||||
let dir = models_dir(s, 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 !matches!(m.role, Role::Embedder | Role::Detector) {
|
||||
return Err(bad(c, "only face detection and recognition models can be installed"));
|
||||
}
|
||||
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 moved = verified && std::fs::rename(src, &part).is_ok();
|
||||
let placed = if moved {
|
||||
Ok(())
|
||||
} else {
|
||||
std::fs::copy(src, &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) = placed {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(e);
|
||||
}
|
||||
// Validate the candidate before replacing a working installation.
|
||||
let test = match self_test(&part, &m) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
let _ = std::fs::remove_dir(&home);
|
||||
return Err(bad(c, e));
|
||||
}
|
||||
};
|
||||
if let Err(e) = std::fs::rename(&part, &final_path) {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(fail("could not finish the copy", 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, "selfTest": test});
|
||||
write_atomic(&home.join("installed.json"), &serde_json::to_vec_pretty(&accepted).map_err(|e| fail("record", e))?)?;
|
||||
// a file that was downloaded and waited in the staging folder has served: the installed copy is the model now
|
||||
if src.parent() == Some(dir.join(face_download::STAGING).as_path()) {
|
||||
let _ = std::fs::remove_file(src);
|
||||
s.face_downloads.staged_file_used(src);
|
||||
}
|
||||
// the newest model is the one in use (the earlier one stays installed, and `faces.models.select` switches back)
|
||||
if activate && m.role == Role::Embedder {
|
||||
let mut st = read_settings(&dir);
|
||||
st.embedder = Some(m.id.clone());
|
||||
st.enabled = true;
|
||||
write_settings(&dir, &st)?;
|
||||
}
|
||||
let now = read_settings(&dir);
|
||||
let selected = now.embedder.as_deref() == Some(m.id.as_str());
|
||||
Ok(json!({"installed": row(s, &m, true, selected, &accepted), "active": {"embedder": now.embedder, "enabled": now.enabled}}))
|
||||
}
|
||||
|
||||
/// What an error says, without the command it came from (for a line shown to the user).
|
||||
fn plain(e: &EngineError) -> String {
|
||||
match e {
|
||||
EngineError::BadParams { msg, .. } => msg.clone(),
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Install every download that has arrived: each was accepted when it was started, so it is installed, chosen and
|
||||
/// switched on without another question. Called by `faces.models.downloads` and by every frame's `faces.pump`.
|
||||
pub(super) fn finish_downloads(s: &mut Session) {
|
||||
const C: &str = "faces.models.download";
|
||||
let extra = catalog_manifests(s);
|
||||
for (id, path, sha) in s.face_downloads.finished() {
|
||||
let outcome = inspect_file(&path.to_string_lossy(), C, Some(&sha), &extra).and_then(|ins| install_file(s, C, &path, ins, true, true, true));
|
||||
match outcome {
|
||||
Ok(_) => s.face_downloads.set_outcome(&id, State::Installed),
|
||||
Err(e) => {
|
||||
// a file that cannot be installed is of no use to anyone: throw it away, say why
|
||||
let _ = std::fs::remove_file(&path);
|
||||
s.face_downloads.set_outcome(&id, State::Failed(plain(&e)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `faces.models.download {id, acknowledged: true}`: fetch a model LightCraft knows an address for, in the background;
|
||||
/// when it has arrived and checked out it is installed.
|
||||
fn download(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.download";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
refresh_catalog(s);
|
||||
let spec = download_spec(s, id).ok_or_else(|| bad(C, "LightCraft has no download for that model: get the file from its page, then add it"))?;
|
||||
if cfg!(target_arch = "wasm32") {
|
||||
return Err(bad(C, "the web build downloads no models"));
|
||||
}
|
||||
if p.get("acknowledged").and_then(Value::as_bool) != Some(true) {
|
||||
return Err(bad(
|
||||
C,
|
||||
"the licence has not been accepted: show the user the model's terms (see `faces.models.list`) and pass `acknowledged: true` once they agree",
|
||||
));
|
||||
}
|
||||
if installed_models(&dir).iter().any(|i| i.manifest.id == spec.id) {
|
||||
return Err(bad(C, "that model is already installed"));
|
||||
}
|
||||
let host = spec.host().to_string();
|
||||
s.face_downloads.start(spec, &dir).map_err(|e| bad(C, e))?;
|
||||
Ok(json!({"started": id, "from": host}))
|
||||
}
|
||||
|
||||
fn download_row(s: &Session, id: &str, state: &State) -> Value {
|
||||
let from = download_spec(s, id).map(|d| d.host().to_string());
|
||||
match state {
|
||||
State::Running { bytes, total } => json!({"id": id, "state": "running", "bytes": bytes, "total": total, "from": from}),
|
||||
State::Done { path, .. } => json!({"id": id, "state": "done", "path": path.display().to_string(), "from": from}),
|
||||
State::Installed => json!({"id": id, "state": "installed", "from": from}),
|
||||
State::Failed(why) => json!({"id": id, "state": "failed", "error": why, "from": from}),
|
||||
State::Cancelled => json!({"id": id, "state": "cancelled", "from": from}),
|
||||
}
|
||||
}
|
||||
|
||||
/// `faces.models.downloads`: where each download stands (a finished one is installed first).
|
||||
fn downloads(s: &mut Session, _: &Value) -> Result<Value> {
|
||||
finish_downloads(s);
|
||||
let all: Vec<Value> = s.face_downloads.snapshot().iter().map(|(id, st)| download_row(s, id, st)).collect();
|
||||
Ok(json!({"running": s.face_downloads.running(), "downloads": all}))
|
||||
}
|
||||
|
||||
/// `faces.models.downloadCancel {id}`: stop a download, or clear a finished one (a staged file is deleted).
|
||||
fn download_cancel(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.downloadCancel";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
Ok(json!({"discarded": s.face_downloads.discard(id)}))
|
||||
}
|
||||
|
||||
/// Load the model and run its self-test: `Ok(null)` in a build without the recognition runtime (nothing can
|
||||
/// be checked), `Ok(result)` when it passes, `Err(reason)` when it does not load or fails a check.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
fn self_test(path: &Path, m: &ModelManifest) -> std::result::Result<Value, String> {
|
||||
if m.role == Role::Detector {
|
||||
let bytes = read_capped(path, 8 << 20).ok_or("the detector cannot be read within its 8 MiB limit")?;
|
||||
let detector = lightcraft_faces::yunet::Detector::new(&bytes).map_err(|e| e.to_string())?;
|
||||
if Some(detector.input_side()) != usize::try_from(m.input.width).ok()
|
||||
|| m.input.width != m.input.height
|
||||
|| !matches!(&m.output, manifest::OutputSpec::Detector { decoder } if decoder == "yunet-v2")
|
||||
{
|
||||
return Err("unsupported detector input or output decoder".into());
|
||||
}
|
||||
detector.detect(&[0u8; 12], 2, 2, &lightcraft_faces::yunet::Options::default()).map_err(|e| e.to_string())?;
|
||||
return Ok(json!({"ok": true}));
|
||||
}
|
||||
let embedder = lightcraft_faces::runtime::Embedder::load(path, m).map_err(|e| e.to_string())?;
|
||||
let t = embedder.self_test();
|
||||
if !t.ok {
|
||||
let failed: Vec<&str> = t.checks.iter().filter(|(_, ok)| !*ok).map(|(c, _)| c.as_str()).collect();
|
||||
return Err(format!("the model failed its self-test: it does not {}", failed.join(", and does not ")));
|
||||
}
|
||||
serde_json::to_value(&t).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
fn self_test(_: &Path, _: &ModelManifest) -> std::result::Result<Value, String> {
|
||||
Ok(Value::Null)
|
||||
}
|
||||
|
||||
/// `faces.models.test {id}`: run an installed model's self-test again.
|
||||
fn test(s: &mut Session, p: &Value) -> Result<Value> {
|
||||
const C: &str = "faces.models.test";
|
||||
let id = str_param(p, "id").ok_or_else(|| bad(C, "missing `id`"))?;
|
||||
let dir = models_dir(s, C)?;
|
||||
if !cfg!(not(target_arch = "wasm32")) {
|
||||
return Err(bad(C, "this build cannot run face recognition models"));
|
||||
}
|
||||
let found = installed_models(&dir).into_iter().find(|i| i.manifest.id == id).ok_or_else(|| bad(C, "that model is not installed"))?;
|
||||
if found.manifest.role != Role::Embedder {
|
||||
return Err(bad(C, "only recognition models have a self-test"));
|
||||
}
|
||||
let result = self_test(&dir.join(id).join("model.onnx"), &found.manifest);
|
||||
let (ok, detail) = match &result {
|
||||
Ok(v) => (true, v.clone()),
|
||||
Err(e) => (false, json!(e)),
|
||||
};
|
||||
// keep the latest outcome with the model's record
|
||||
let mut record = found.accepted;
|
||||
if let Some(o) = record.as_object_mut() {
|
||||
o.insert("selfTest".into(), if ok { detail.clone() } else { json!({"ok": false, "error": detail}) });
|
||||
write_atomic(&dir.join(id).join("installed.json"), &serde_json::to_vec_pretty(&record).map_err(|e| fail("record", e))?)?;
|
||||
}
|
||||
Ok(json!({"id": id, "ok": ok, "result": detail}))
|
||||
}
|
||||
|
||||
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) {
|
||||
// another installed recogniser takes over, so recognition keeps working; with none left nothing is chosen
|
||||
st.embedder = installed_models(&dir).into_iter().find(|i| i.manifest.id != id && i.manifest.role == Role::Embedder).map(|i| i.manifest.id);
|
||||
write_settings(&dir, &st)?;
|
||||
}
|
||||
Ok(json!({"removed": id, "embedder": st.embedder}))
|
||||
}
|
||||
|
||||
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)?;
|
||||
// the background scan notices at once, not within the second it trusts its last look
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
{
|
||||
s.faces.enabled_seen = None;
|
||||
}
|
||||
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, activate?: true} → {installed, active: {embedder, enabled}} — 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. Unless `activate` is false the model becomes the one in use and recognition is switched on; an earlier model stays installed", always, install),
|
||||
cmd!(query "faces.models.download", "Download Face Model", [], None, "{id, acknowledged: true} → {started, from} — fetch a recognition model LightCraft has a pinned address for (see `downloadHost` in the list), in the background over pure-Rust https (lightcraft-fetch). `acknowledged` must be true: the user has been shown the model's terms and accepted them. It is checked against its size and SHA-256, then installed, chosen and switched on by itself; `faces.models.downloads` shows how far it is", always, download),
|
||||
cmd!(query "faces.models.downloads", "Face Model Downloads", [], None, "{} → {running, downloads: [{id, state: running | done | installed | failed | cancelled, bytes, total, error, from}]} — also installs any download that has arrived; `installed` stays listed until `faces.models.downloadCancel` clears it", always, downloads),
|
||||
cmd!(query "faces.models.downloadCancel", "Cancel Face Model Download", [], None, "{id} → {discarded} — stop a download, or delete a finished one that was not installed", always, download_cancel),
|
||||
cmd!(query "faces.models.test", "Test Face Model", [], None, "{id} → {ok, result: {loadMs, embedMs, dimension, checks}} — load an installed recognition model and check it gives sensible faces; needs the recognition runtime", always, test),
|
||||
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),
|
||||
]
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -119,6 +119,14 @@ fn album_param(p: &Value, key: &str, c: &str) -> Result<AlbumId> {
|
||||
p.get(key).and_then(Value::as_u64).map(AlbumId).ok_or_else(|| bad(c, format!("missing album `{key}`")))
|
||||
}
|
||||
|
||||
/// An optional album id: absent or `null` is `None`, anything but a number is an error.
|
||||
fn opt_album_param(p: &Value, key: &str, c: &str) -> Result<Option<AlbumId>> {
|
||||
match p.get(key) {
|
||||
None | Some(Value::Null) => Ok(None),
|
||||
Some(v) => v.as_u64().map(|n| Some(AlbumId(n))).ok_or_else(|| bad(c, format!("`{key}` must be an album id or null"))),
|
||||
}
|
||||
}
|
||||
|
||||
fn strs(p: &Value, key: &str) -> Vec<String> {
|
||||
p.get(key).and_then(Value::as_array).map(|a| a.iter().filter_map(Value::as_str).map(str::to_string).collect()).unwrap_or_default()
|
||||
}
|
||||
@@ -784,7 +792,7 @@ pub fn specs() -> Vec<CommandSpec> {
|
||||
let cover = photos.first().copied();
|
||||
s.commit(
|
||||
if folder { "New Folder" } else { "New Album" },
|
||||
Op::AddAlbum { album: Album { id, name, parent, folder, photos, cover, smart: None, quick: false } },
|
||||
Op::AddAlbum { album: Album { id, name, parent, folder, photos, cover, smart: None, quick: false, order: None } },
|
||||
)?;
|
||||
Ok(json!({"id": id.0}))
|
||||
}),
|
||||
@@ -880,6 +888,80 @@ pub fn specs() -> Vec<CommandSpec> {
|
||||
s.commit("Move Album", Op::MoveAlbum { id, parent })?;
|
||||
ok()
|
||||
}),
|
||||
cmd!(
|
||||
"album.reorder",
|
||||
"Reorder Album",
|
||||
[],
|
||||
None,
|
||||
"{id, parent?: folderId|null, before?: albumId|null} — places the album (or folder) in `parent` (default: where it is; null: the top level) before the sibling `before` of the same kind, or last of its kind; the folder is then ordered by hand. One undo step. → {changed}: how many edits it took (0: already there)",
|
||||
always,
|
||||
|s, p| {
|
||||
let id = album_param(p, "id", "album.reorder")?;
|
||||
let me = s.catalog.album(id).ok_or_else(|| bad("album.reorder", "no such album"))?;
|
||||
let parent = if p.get("parent").is_some() { opt_album_param(p, "parent", "album.reorder")? } else { me.parent };
|
||||
let before = opt_album_param(p, "before", "album.reorder")?;
|
||||
if before == Some(id) {
|
||||
// before itself: where it is
|
||||
return Ok(json!({"changed": 0}));
|
||||
}
|
||||
let mut order: Vec<AlbumId> = s.catalog.album_children(parent).into_iter().filter(|a| a.id != id).map(|a| a.id).collect();
|
||||
let at = match before {
|
||||
Some(b) => {
|
||||
let at = order.iter().position(|x| *x == b).ok_or_else(|| bad("album.reorder", "`before` is not inside that folder"))?;
|
||||
if s.catalog.album(b).is_none_or(|x| x.folder != me.folder) {
|
||||
return Err(bad("album.reorder", "folders are placed among folders and albums among albums"));
|
||||
}
|
||||
at
|
||||
}
|
||||
None if me.folder => order.iter().take_while(|x| s.catalog.album(**x).is_some_and(|a| a.folder)).count(),
|
||||
None => order.len(),
|
||||
};
|
||||
order.insert(at.min(order.len()), id);
|
||||
let moved = me.parent != parent;
|
||||
let mut ops = Vec::new();
|
||||
if moved {
|
||||
ops.push(Op::MoveAlbum { id, parent });
|
||||
}
|
||||
for (i, aid) in order.iter().enumerate() {
|
||||
let want = u32::try_from(i).ok();
|
||||
// a moved album starts without a place in its new folder
|
||||
let have = if moved && *aid == id { None } else { s.catalog.album(*aid).and_then(|a| a.order) };
|
||||
if have != want {
|
||||
ops.push(Op::SetAlbumOrder { id: *aid, order: want });
|
||||
}
|
||||
}
|
||||
if ops.is_empty() {
|
||||
return Ok(json!({"changed": 0}));
|
||||
}
|
||||
let changed = ops.len();
|
||||
s.commit("Reorder Album", Op::Batch { ops })?;
|
||||
Ok(json!({"changed": changed}))
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
"album.sort",
|
||||
"Sort Albums by Name",
|
||||
[],
|
||||
None,
|
||||
"{parent?: folderId|null} — drops the hand order of the albums inside the folder (default: the top level), so they are listed by name → {changed}: how many albums lost their place",
|
||||
always,
|
||||
|s, p| {
|
||||
let parent = opt_album_param(p, "parent", "album.sort")?;
|
||||
let ops: Vec<Op> = s
|
||||
.catalog
|
||||
.album_children(parent)
|
||||
.into_iter()
|
||||
.filter(|a| a.order.is_some())
|
||||
.map(|a| Op::SetAlbumOrder { id: a.id, order: None })
|
||||
.collect();
|
||||
if ops.is_empty() {
|
||||
return Ok(json!({"changed": 0}));
|
||||
}
|
||||
let changed = ops.len();
|
||||
s.commit("Sort Albums by Name", Op::Batch { ops })?;
|
||||
Ok(json!({"changed": changed}))
|
||||
}
|
||||
),
|
||||
cmd!("album.addPhotos", "Add to Album", ["Photo"], None, "{id: albumId, ids?: [photoIds]} (default: selection)", has_selection, |s, p| {
|
||||
let id = album_param(p, "id", "album.addPhotos")?;
|
||||
let targets = ids_param(p).unwrap_or_else(|| s.targets(&Value::Null));
|
||||
|
||||
@@ -432,6 +432,7 @@ pub fn specs() -> Vec<CommandSpec> {
|
||||
"available": crate::segment::Segmenter::AVAILABLE,
|
||||
"installed": g.installed(),
|
||||
"dir": g.dir.as_ref().map(|d| d.display().to_string()),
|
||||
"remote": g.remote_endpoint(),
|
||||
"loaded": g.loaded(),
|
||||
"busy": g.busy(),
|
||||
"analyzing": g.analyzing(),
|
||||
|
||||
@@ -10,9 +10,13 @@ mod color;
|
||||
pub(crate) mod convert;
|
||||
mod cull;
|
||||
pub mod curves;
|
||||
pub(crate) mod denoise;
|
||||
mod develop;
|
||||
mod edit;
|
||||
mod export;
|
||||
pub(crate) mod face_detect;
|
||||
mod face_models;
|
||||
mod face_recognize;
|
||||
pub mod filters;
|
||||
pub mod keywords;
|
||||
pub mod library;
|
||||
@@ -140,6 +144,10 @@ pub fn command_specs() -> &'static [CommandSpec] {
|
||||
v.extend(missing::specs());
|
||||
v.extend(metadata::specs());
|
||||
v.extend(filters::specs());
|
||||
v.extend(denoise::specs());
|
||||
v.extend(face_models::specs());
|
||||
v.extend(face_detect::specs());
|
||||
v.extend(face_recognize::specs());
|
||||
v
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
//! Where LightCraft keeps per-user data on this machine: the config folder (app settings, camera
|
||||
//! profiles, face and denoise 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- or denoise-models folder until a host asks for one ([`Session::with_default_face_models`],
|
||||
//! [`Session::with_default_denoise_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")))
|
||||
}
|
||||
|
||||
/// Where opt-in denoise models are kept: `$LIGHTCRAFT_DENOISE_MODELS` if set, else `<config>/denoise-models`.
|
||||
pub fn default_denoise_models_dir() -> Option<PathBuf> {
|
||||
std::env::var_os("LIGHTCRAFT_DENOISE_MODELS")
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| config_dir().map(|d| d.join("denoise-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
|
||||
}
|
||||
|
||||
/// Keep denoise models in the shared default folder ([`default_denoise_models_dir`]).
|
||||
pub fn with_default_denoise_models(mut self) -> Self {
|
||||
self.set_denoise_models_dir(default_denoise_models_dir());
|
||||
self
|
||||
}
|
||||
|
||||
/// Point AI denoise at `dir` (or at no folder), dropping what was learned about the old one.
|
||||
pub fn set_denoise_models_dir(&mut self, dir: Option<PathBuf>) {
|
||||
self.denoise.models_dir = dir;
|
||||
self.denoise.touch();
|
||||
}
|
||||
}
|
||||
@@ -62,7 +62,17 @@ pub fn load(s: &mut Session) {
|
||||
}
|
||||
let folder = s.catalog.alloc_album_id();
|
||||
ops.push(Op::AddAlbum {
|
||||
album: Album { id: folder, name: "Travel 2026".into(), parent: None, folder: true, photos: vec![], cover: None, smart: None, quick: false },
|
||||
album: Album {
|
||||
id: folder,
|
||||
name: "Travel 2026".into(),
|
||||
parent: None,
|
||||
folder: true,
|
||||
photos: vec![],
|
||||
cover: None,
|
||||
smart: None,
|
||||
quick: false,
|
||||
order: None,
|
||||
},
|
||||
});
|
||||
let by_kw = |kw: &str| -> Vec<_> { scenes.iter().zip(&ids).filter(|(sc, _)| sc.meta.keywords.contains(&kw)).map(|(_, id)| *id).collect() };
|
||||
for (name, parent, photos) in [
|
||||
@@ -75,7 +85,9 @@ pub fn load(s: &mut Session) {
|
||||
] {
|
||||
let id = s.catalog.alloc_album_id();
|
||||
let cover = photos.first().copied();
|
||||
ops.push(Op::AddAlbum { album: Album { id, name: name.into(), parent, folder: false, photos, cover, smart: None, quick: false } });
|
||||
ops.push(Op::AddAlbum {
|
||||
album: Album { id, name: name.into(), parent, folder: false, photos, cover, smart: None, quick: false, order: None },
|
||||
});
|
||||
}
|
||||
for op in ops {
|
||||
let _ = s.catalog.apply(op);
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,372 @@
|
||||
//! Downloading a face model the user asked for (Settings ▸ Faces ▸ Download).
|
||||
//!
|
||||
//! The transfer is [`lightcraft_fetch`] (pure-Rust HTTPS with resume, size and SHA-256 checks), on a background thread,
|
||||
//! so the UI never waits for the network. Only a model listed in [`lightcraft_faces::known::download`] can be fetched:
|
||||
//! the address is a pinned constant, never something typed or passed in, and no token is ever sent with it. What arrives
|
||||
//! matches the model's recorded size and SHA-256 or it is thrown away; it then waits in a staging folder until the model
|
||||
//! is installed (`faces.models.install`, or by itself once the user has accepted the model's terms in the download).
|
||||
//!
|
||||
//! Failure is a message, never a panic: no network, a stalled or truncated transfer, a full disk, a file that is not the
|
||||
//! one expected, and a panic inside the thread itself all end as [`State::Failed`].
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex, PoisonError};
|
||||
use std::time::Duration;
|
||||
|
||||
use lightcraft_faces::known::Download;
|
||||
|
||||
/// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it
|
||||
/// can never be a model's own folder (ids start with a letter or digit).
|
||||
pub const STAGING: &str = ".downloads";
|
||||
|
||||
/// Where a download stands.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum State {
|
||||
/// Fetching (`bytes` of `total` so far); the file is checked as the last step.
|
||||
Running {
|
||||
bytes: u64,
|
||||
total: u64,
|
||||
},
|
||||
/// Fetched and verified (its SHA-256 is `sha256`), waiting in the staging folder at `path` to be installed.
|
||||
Done {
|
||||
path: PathBuf,
|
||||
sha256: String,
|
||||
},
|
||||
/// Installed and in use: nothing more to do but tell the user.
|
||||
Installed,
|
||||
Failed(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
state: Mutex<State>,
|
||||
cancel: AtomicBool,
|
||||
}
|
||||
|
||||
impl Shared {
|
||||
fn new(total: u64) -> Self {
|
||||
Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) }
|
||||
}
|
||||
|
||||
fn set(&self, s: State) {
|
||||
*self.state.lock().unwrap_or_else(PoisonError::into_inner) = s;
|
||||
}
|
||||
|
||||
fn get(&self) -> State {
|
||||
self.state.lock().unwrap_or_else(PoisonError::into_inner).clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// The downloads of one session, by model id.
|
||||
#[derive(Default)]
|
||||
pub struct Downloads {
|
||||
jobs: HashMap<String, Arc<Shared>>,
|
||||
}
|
||||
|
||||
impl Downloads {
|
||||
/// Start fetching `spec` into `<models dir>/.downloads/`. Different models can download side by side; the same
|
||||
/// one is started once.
|
||||
pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> {
|
||||
if self.jobs.get(spec.id.as_str()).is_some_and(|j| matches!(j.get(), State::Running { .. })) {
|
||||
return Err("that model is already downloading".into());
|
||||
}
|
||||
let staging = models_dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?;
|
||||
let shared = Arc::new(Shared::new(spec.size_bytes));
|
||||
let worker = shared.clone();
|
||||
let id = spec.id.clone();
|
||||
std::thread::Builder::new()
|
||||
.name("face-model-download".into())
|
||||
.spawn(move || {
|
||||
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, spec.base(), &staging, &worker)));
|
||||
if outcome.is_err() {
|
||||
worker.set(State::Failed("the download stopped unexpectedly".into()));
|
||||
}
|
||||
})
|
||||
.map_err(|e| format!("could not start the download: {e}"))?;
|
||||
self.jobs.insert(id, shared);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A download that has arrived and been checked, as if the transfer had just ended (tests of installing).
|
||||
#[cfg(test)]
|
||||
pub(crate) fn arrived(&mut self, id: &str, path: PathBuf, sha256: &str) {
|
||||
let j = Shared::new(0);
|
||||
j.set(State::Done { path, sha256: sha256.to_string() });
|
||||
self.jobs.insert(id.to_string(), Arc::new(j));
|
||||
}
|
||||
|
||||
/// Every download this session has started and not discarded, by model id.
|
||||
pub fn snapshot(&self) -> Vec<(String, State)> {
|
||||
let mut all: Vec<_> = self.jobs.iter().map(|(id, j)| (id.clone(), j.get())).filter(|(_, s)| *s != State::Cancelled).collect();
|
||||
all.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
all
|
||||
}
|
||||
|
||||
pub fn running(&self) -> bool {
|
||||
self.jobs.values().any(|j| matches!(j.get(), State::Running { .. }))
|
||||
}
|
||||
|
||||
/// Stop a running download, or discard a finished one (its staged file is deleted). `true` if there was one.
|
||||
pub fn discard(&mut self, id: &str) -> bool {
|
||||
let Some(job) = self.jobs.get(id) else { return false };
|
||||
job.cancel.store(true, Ordering::Relaxed);
|
||||
match job.get() {
|
||||
// the thread notices, deletes the partial file and marks itself cancelled
|
||||
State::Running { .. } => {}
|
||||
State::Done { path, .. } => {
|
||||
let _ = std::fs::remove_file(path);
|
||||
self.jobs.remove(id);
|
||||
}
|
||||
State::Installed | State::Failed(_) | State::Cancelled => {
|
||||
self.jobs.remove(id);
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// The staged file at `path` has been installed (its copy is the model now) or is no longer wanted.
|
||||
pub fn staged_file_used(&mut self, path: &Path) {
|
||||
self.jobs.retain(|_, j| !matches!(j.get(), State::Done { path: p, .. } if p == path));
|
||||
}
|
||||
|
||||
/// Downloads that have arrived and been checked, ready to install: (model id, staged file, its SHA-256).
|
||||
pub fn finished(&self) -> Vec<(String, PathBuf, String)> {
|
||||
let mut out: Vec<_> = self
|
||||
.jobs
|
||||
.iter()
|
||||
.filter_map(|(id, j)| match j.get() {
|
||||
State::Done { path, sha256 } => Some((id.clone(), path, sha256)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
out.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
out
|
||||
}
|
||||
|
||||
/// Record how installing a finished download went (`Installed` or `Failed`).
|
||||
pub fn set_outcome(&mut self, id: &str, state: State) {
|
||||
match self.jobs.get(id) {
|
||||
Some(j) => j.set(state),
|
||||
None => {
|
||||
let j = Shared::new(0);
|
||||
j.set(state);
|
||||
self.jobs.insert(id.to_string(), Arc::new(j));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Downloads {
|
||||
/// Quitting while a download runs stops it instead of leaving it behind.
|
||||
fn drop(&mut self) {
|
||||
for job in self.jobs.values() {
|
||||
job.cancel.store(true, Ordering::Relaxed);
|
||||
}
|
||||
let started = std::time::Instant::now();
|
||||
while self.running() && started.elapsed() < Duration::from_secs(1) {
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What went wrong, in words the user can act on.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
fn explain(e: &lightcraft_fetch::DownloadError) -> String {
|
||||
use lightcraft_fetch::DownloadError;
|
||||
match e {
|
||||
DownloadError::AllMirrorsFailed { errors, .. } => {
|
||||
let detail: String = errors.join("; ").chars().take(300).collect();
|
||||
format!("could not download the model ({detail}): check the internet connection, or get the file from its page and add it yourself")
|
||||
}
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch `spec` from `mirror` (the address without the file name) into `staging`, check it and leave it there.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
fn fetch(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) -> Result<(), lightcraft_fetch::DownloadError> {
|
||||
let file = lightcraft_fetch::FileSpec { name: &spec.file_name, size: Some(spec.size_bytes), sha256: Some(&spec.sha256), max: spec.size_bytes };
|
||||
lightcraft_fetch::download(&[file], &[mirror.to_string()], staging, &lightcraft_fetch::Options::default(), &shared.cancel, &mut |p| {
|
||||
shared.set(State::Running { bytes: p.done.min(spec.size_bytes), total: spec.size_bytes });
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
fn fetch(_: &Download, _: &str, _: &Path, _: &Shared) -> Result<(), String> {
|
||||
Err("the web build downloads no models".into())
|
||||
}
|
||||
|
||||
/// The whole job: reuse a verified earlier copy, else fetch, check and stage. Always ends in a final [`State`].
|
||||
fn run(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) {
|
||||
let target = staging.join(&spec.file_name);
|
||||
let part = staging.join(format!("{}.part", spec.file_name));
|
||||
match fetch(spec, mirror, staging, shared) {
|
||||
Ok(()) => shared.set(State::Done { path: target, sha256: spec.sha256.to_string() }),
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
Err(lightcraft_fetch::DownloadError::Cancelled) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
shared.set(State::Cancelled);
|
||||
}
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
shared.set(State::Failed(explain(&e)));
|
||||
}
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
shared.set(State::Failed(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, not(target_arch = "wasm32")))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use lightcraft_faces::hash::sha256_hex;
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::net::TcpListener;
|
||||
|
||||
fn scratch(name: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
fn spec_for(bytes: &[u8]) -> Download {
|
||||
Download {
|
||||
id: "test-model".into(),
|
||||
url: "https://example.invalid/m/test-model.onnx".into(),
|
||||
file_name: "test-model.onnx".into(),
|
||||
size_bytes: bytes.len() as u64,
|
||||
sha256: sha256_hex(bytes),
|
||||
}
|
||||
}
|
||||
|
||||
/// A local server that answers every request with `status` and `body` (the whole file, no ranges).
|
||||
fn serve(status: &'static str, body: Vec<u8>) -> String {
|
||||
let l = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let base = format!("http://{}/m", l.local_addr().unwrap());
|
||||
std::thread::spawn(move || {
|
||||
for mut s in l.incoming().flatten() {
|
||||
let mut r = BufReader::new(s.try_clone().unwrap());
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = s.write_all(format!("HTTP/1.1 {status}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes());
|
||||
let _ = s.write_all(&body);
|
||||
}
|
||||
});
|
||||
base
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fetches_checks_and_stages_a_file() {
|
||||
let dir = scratch("ok");
|
||||
let bytes: Vec<u8> = (0..200_000u32).map(|i| (i % 251) as u8).collect();
|
||||
let (staging, spec) = (dir.join(STAGING), spec_for(&bytes));
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let mirror = serve("200 OK", bytes.clone());
|
||||
let shared = Shared::new(spec.size_bytes);
|
||||
run(&spec, &mirror, &staging, &shared);
|
||||
let State::Done { path, sha256 } = shared.get() else { panic!("{:?}", shared.get()) };
|
||||
assert_eq!((std::fs::read(&path).unwrap(), sha256), (bytes, spec.sha256.to_string()));
|
||||
assert!(!staging.join("test-model.onnx.part").exists());
|
||||
// asked again with the copy still staged: no fetch needed (nothing listens there)
|
||||
let again = Shared::new(spec.size_bytes);
|
||||
run(&spec, "http://127.0.0.1:1/m", &staging, &again);
|
||||
assert!(matches!(again.get(), State::Done { .. }), "{:?}", again.get());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_file_that_does_not_match_is_thrown_away() {
|
||||
let dir = scratch("bad");
|
||||
let staging = dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
// the right size and the wrong content
|
||||
let spec = spec_for(&[7u8; 5000]);
|
||||
let shared = Shared::new(spec.size_bytes);
|
||||
run(&spec, &serve("200 OK", vec![8u8; 5000]), &staging, &shared);
|
||||
assert!(matches!(shared.get(), State::Failed(_)), "{:?}", shared.get());
|
||||
// nothing is left behind: not the file, not a partial one
|
||||
assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unreachable_or_refusing_server_is_a_message() {
|
||||
let dir = scratch("down");
|
||||
let staging = dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let spec = spec_for(b"abc");
|
||||
for mirror in ["http://127.0.0.1:1/m".to_string(), serve("404 Not Found", b"not found".to_vec())] {
|
||||
let shared = Shared::new(3);
|
||||
run(&spec, &mirror, &staging, &shared);
|
||||
let State::Failed(why) = shared.get() else { panic!("{:?}", shared.get()) };
|
||||
assert!(why.contains("internet connection"), "{why}");
|
||||
}
|
||||
assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_cancelled_download_leaves_nothing() {
|
||||
let dir = scratch("cancel");
|
||||
let staging = dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let spec = spec_for(b"abc");
|
||||
let shared = Shared::new(3);
|
||||
shared.cancel.store(true, Ordering::Relaxed);
|
||||
run(&spec, &serve("200 OK", b"abc".to_vec()), &staging, &shared);
|
||||
assert_eq!(shared.get(), State::Cancelled);
|
||||
assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dropping_the_session_stops_a_running_download() {
|
||||
let shared = Arc::new(Shared::new(10));
|
||||
// a stand-in for the download thread: notices the cancel flag the way the downloader does
|
||||
let worker = shared.clone();
|
||||
let handle = std::thread::spawn(move || {
|
||||
while !worker.cancel.load(Ordering::Relaxed) {
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
worker.set(State::Cancelled);
|
||||
});
|
||||
let mut d = Downloads::default();
|
||||
d.jobs.insert("m".into(), shared.clone());
|
||||
assert!(d.running());
|
||||
let started = std::time::Instant::now();
|
||||
drop(d);
|
||||
assert!(started.elapsed() < Duration::from_millis(900), "quitting must not wait for a transfer to finish");
|
||||
assert_eq!(shared.get(), State::Cancelled);
|
||||
handle.join().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discarding_a_finished_download_deletes_its_file() {
|
||||
let dir = scratch("discard");
|
||||
let file = dir.join("staged.onnx");
|
||||
std::fs::write(&file, b"x").unwrap();
|
||||
let mut d = Downloads::default();
|
||||
let shared = Arc::new(Shared::new(1));
|
||||
shared.set(State::Done { path: file.clone(), sha256: String::new() });
|
||||
d.jobs.insert("m".into(), shared);
|
||||
assert_eq!(d.snapshot().len(), 1);
|
||||
assert!(!d.running());
|
||||
assert!(d.discard("m"));
|
||||
assert!(!file.exists());
|
||||
assert!(d.snapshot().is_empty());
|
||||
assert!(!d.discard("m"), "nothing left to discard");
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,455 @@
|
||||
//! The face index: one embedding per face region, kept in memory and cached beside the library.
|
||||
//!
|
||||
//! Embedding a face costs a model run (tens to hundreds of milliseconds), so each region is embedded once and
|
||||
//! remembered, keyed by its photo and its box. The cache file (`face-embeddings-<model id>.bin` in the library folder,
|
||||
//! one per model, so switching back to an earlier model does not start over) belongs to one model: its header names the
|
||||
//! model and its file hash, and a file made by another model, an older format, or a damaged one is simply ignored and
|
||||
//! rebuilt. Embeddings are kept out of the catalog on purpose: they are large, rebuildable, and not comparable between
|
||||
//! models.
|
||||
//!
|
||||
//! Each face also keeps the box to cut its picture from (its "view"): the detector's box when the detector found the same
|
||||
//! face, so every face is shown equally close however loosely or tightly its own region was drawn (a Lightroom box, an
|
||||
//! older tool's box), else the region's own box.
|
||||
//!
|
||||
//! Layout (little-endian): `"LCFE"`, version `2`, tag length `u16`, the tag (`<model id>@<sha256>`), dimension `u32`,
|
||||
//! then records of `photo u64`, four `i32` (the region's box in ten-thousandths), four `i32` (the view box likewise), and
|
||||
//! `dim` `f32`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use lightcraft_geom::Rect;
|
||||
|
||||
const MAGIC: &[u8; 4] = b"LCFE";
|
||||
const VERSION: u8 = 2;
|
||||
/// Largest cache file we read.
|
||||
const MAX_FILE: u64 = 1 << 30;
|
||||
|
||||
pub(crate) type Key = (u64, [i32; 4]);
|
||||
|
||||
/// One face's record: its embedding and its view box (quantized like a key).
|
||||
type Entry = (Arc<[f32]>, [i32; 4]);
|
||||
|
||||
/// A quantized box as a rectangle again.
|
||||
fn rect_of(k: [i32; 4]) -> Rect {
|
||||
let f = |v: i32| f64::from(v) / 10_000.0;
|
||||
Rect { x0: f(k[0]), y0: f(k[1]), x1: f(k[2]), y1: f(k[3]) }
|
||||
}
|
||||
|
||||
/// A region's box as the integer key it is cached under.
|
||||
pub(crate) fn region_key(r: &Rect) -> [i32; 4] {
|
||||
let q = |v: f64| if v.is_finite() { (v.clamp(-1e3, 1e3) * 10_000.0).round() as i32 } else { 0 };
|
||||
[q(r.x0), q(r.y0), q(r.x1), q(r.y1)]
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct Index {
|
||||
/// `<model id>@<sha256>` of the model the embeddings belong to.
|
||||
pub tag: String,
|
||||
pub dim: usize,
|
||||
entries: HashMap<Key, Entry>,
|
||||
unsaved: Vec<(Key, Entry)>,
|
||||
/// Faces that could not be embedded this session (so they are not retried on every call).
|
||||
skipped: std::collections::HashSet<Key>,
|
||||
/// The cache file on disk belongs to this model, so new records can be appended to it.
|
||||
file_matches: bool,
|
||||
pub path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl Index {
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries.len()
|
||||
}
|
||||
|
||||
pub fn get(&self, photo: u64, rect: &Rect) -> Option<&Arc<[f32]>> {
|
||||
self.entries.get(&(photo, region_key(rect))).map(|e| &e.0)
|
||||
}
|
||||
|
||||
/// The box to cut this face's picture from, once the scan has looked at it.
|
||||
pub fn view(&self, photo: u64, rect: &Rect) -> Option<Rect> {
|
||||
self.entries.get(&(photo, region_key(rect))).map(|e| rect_of(e.1))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn contains(&self, photo: u64, rect: &Rect) -> bool {
|
||||
self.entries.contains_key(&(photo, region_key(rect)))
|
||||
}
|
||||
|
||||
/// Whether this face still has to be embedded.
|
||||
pub fn needs(&self, photo: u64, rect: &Rect) -> bool {
|
||||
let key = (photo, region_key(rect));
|
||||
!self.entries.contains_key(&key) && !self.skipped.contains(&key)
|
||||
}
|
||||
|
||||
pub fn skip(&mut self, photo: u64, rect: &Rect) {
|
||||
self.skipped.insert((photo, region_key(rect)));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn insert(&mut self, photo: u64, rect: &Rect, embedding: Vec<f32>) {
|
||||
self.insert_with_view(photo, rect, embedding, rect);
|
||||
}
|
||||
|
||||
/// Remember a face's embedding and the box to show it by.
|
||||
pub fn insert_with_view(&mut self, photo: u64, rect: &Rect, embedding: Vec<f32>, view: &Rect) {
|
||||
let entry: Entry = (embedding.into(), region_key(view));
|
||||
let key = (photo, region_key(rect));
|
||||
self.entries.insert(key, entry.clone());
|
||||
self.unsaved.push((key, entry));
|
||||
}
|
||||
|
||||
/// Whether there are embeddings not yet written to the cache file.
|
||||
pub fn has_unsaved(&self) -> bool {
|
||||
!self.unsaved.is_empty()
|
||||
}
|
||||
|
||||
/// Start over for `tag` (`dim` numbers per face), loading what the cache file at `path` holds for it. What the
|
||||
/// previous model's index had not written yet is written first, so switching models loses nothing.
|
||||
pub fn reset(&mut self, tag: &str, dim: usize, path: Option<PathBuf>) {
|
||||
let _ = self.save();
|
||||
self.tag = tag.to_string();
|
||||
self.dim = dim;
|
||||
self.path = path;
|
||||
self.entries.clear();
|
||||
self.unsaved.clear();
|
||||
self.skipped.clear();
|
||||
self.file_matches = false;
|
||||
if let Some(p) = self.path.clone() {
|
||||
self.load(&p);
|
||||
}
|
||||
}
|
||||
|
||||
fn load(&mut self, path: &Path) {
|
||||
let Ok(meta) = std::fs::metadata(path) else { return };
|
||||
if meta.len() > MAX_FILE {
|
||||
return;
|
||||
}
|
||||
let Ok(bytes) = std::fs::read(path) else { return };
|
||||
let Some((tag_len, rest)) = bytes
|
||||
.strip_prefix(MAGIC)
|
||||
.and_then(|r| r.split_first())
|
||||
.filter(|(v, _)| **v == VERSION)
|
||||
.and_then(|(_, r)| r.split_at_checked(2))
|
||||
.map(|(l, r)| (usize::from(u16::from_le_bytes([l[0], l[1]])), r))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some((tag, rest)) = rest.split_at_checked(tag_len) else { return };
|
||||
let Some((dim, mut rest)) = rest.split_at_checked(4).map(|(d, r)| (u32::from_le_bytes([d[0], d[1], d[2], d[3]]) as usize, r)) else { return };
|
||||
if tag != self.tag.as_bytes() || dim != self.dim || dim == 0 {
|
||||
return;
|
||||
}
|
||||
self.file_matches = true;
|
||||
let record = 8 + 16 + 16 + dim * 4;
|
||||
while let Some((rec, tail)) = rest.split_at_checked(record) {
|
||||
rest = tail;
|
||||
let photo = u64::from_le_bytes([rec[0], rec[1], rec[2], rec[3], rec[4], rec[5], rec[6], rec[7]]);
|
||||
let int = |at: usize| i32::from_le_bytes([rec[at], rec[at + 1], rec[at + 2], rec[at + 3]]);
|
||||
let key = [int(8), int(12), int(16), int(20)];
|
||||
let view = [int(24), int(28), int(32), int(36)];
|
||||
let emb: Vec<f32> = rec[40..].as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect();
|
||||
if emb.iter().all(|v| v.is_finite()) {
|
||||
self.entries.insert((photo, key), (emb.into(), view));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn header(&self) -> Vec<u8> {
|
||||
let tag = self.tag.as_bytes();
|
||||
let mut h = Vec::with_capacity(11 + tag.len());
|
||||
h.extend_from_slice(MAGIC);
|
||||
h.push(VERSION);
|
||||
h.extend_from_slice(&u16::try_from(tag.len()).unwrap_or(0).to_le_bytes());
|
||||
h.extend_from_slice(tag);
|
||||
h.extend_from_slice(&(self.dim as u32).to_le_bytes());
|
||||
h
|
||||
}
|
||||
|
||||
fn record(&self, key: &Key, entry: &Entry) -> Vec<u8> {
|
||||
let mut r = Vec::with_capacity(40 + entry.0.len() * 4);
|
||||
r.extend_from_slice(&key.0.to_le_bytes());
|
||||
key.1.iter().chain(entry.1.iter()).for_each(|k| r.extend_from_slice(&k.to_le_bytes()));
|
||||
entry.0.iter().for_each(|v| r.extend_from_slice(&v.to_le_bytes()));
|
||||
r
|
||||
}
|
||||
|
||||
/// Write what was added since the last save: appended when the file already belongs to this model, else the
|
||||
/// whole index is written fresh (atomically). A failure is returned, the records stay pending.
|
||||
pub fn save(&mut self) -> std::io::Result<()> {
|
||||
let Some(path) = self.path.clone() else {
|
||||
self.unsaved.clear();
|
||||
return Ok(());
|
||||
};
|
||||
if self.unsaved.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
if self.file_matches {
|
||||
let mut f = std::fs::OpenOptions::new().append(true).open(&path)?;
|
||||
let mut buf = Vec::new();
|
||||
for (k, e) in &self.unsaved {
|
||||
buf.extend(self.record(k, e));
|
||||
}
|
||||
f.write_all(&buf)?;
|
||||
} else {
|
||||
let mut buf = self.header();
|
||||
for (k, e) in &self.entries {
|
||||
buf.extend(self.record(k, e));
|
||||
}
|
||||
let part = path.with_extension("part");
|
||||
std::fs::write(&part, &buf)?;
|
||||
std::fs::rename(&part, &path)?;
|
||||
self.file_matches = true;
|
||||
}
|
||||
self.unsaved.clear();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Index {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.save();
|
||||
}
|
||||
}
|
||||
|
||||
const SCAN_MAGIC: &[u8; 4] = b"LCFS";
|
||||
const SCAN_VERSION: u8 = 1;
|
||||
/// The detector a search was made with: a record made by another one is ignored (a better detector may find more).
|
||||
pub(crate) const SCAN_TAG: &str = "yunet-2023mar";
|
||||
/// Largest record file we read.
|
||||
const SCAN_MAX_FILE: u64 = 64 << 20;
|
||||
|
||||
/// The photos the detector has already searched for faces, so each is searched once, however many times LightCraft
|
||||
/// starts. Kept beside the library (`face-scanned.bin`: a header naming the detector, then photo ids, appended to).
|
||||
#[derive(Default)]
|
||||
pub(crate) struct Scanned {
|
||||
set: std::collections::HashSet<u64>,
|
||||
unsaved: Vec<u64>,
|
||||
path: Option<PathBuf>,
|
||||
file_matches: bool,
|
||||
}
|
||||
|
||||
impl Scanned {
|
||||
pub fn path(&self) -> Option<&Path> {
|
||||
self.path.as_deref()
|
||||
}
|
||||
|
||||
pub fn contains(&self, photo: u64) -> bool {
|
||||
self.set.contains(&photo)
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.set.len()
|
||||
}
|
||||
|
||||
pub fn insert(&mut self, photo: u64) {
|
||||
if self.set.insert(photo) {
|
||||
self.unsaved.push(photo);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn has_unsaved(&self) -> bool {
|
||||
!self.unsaved.is_empty()
|
||||
}
|
||||
|
||||
/// Start over with the record at `path` (written first: what was not yet saved is not lost).
|
||||
pub fn reset(&mut self, path: Option<PathBuf>) {
|
||||
let _ = self.save();
|
||||
self.set.clear();
|
||||
self.unsaved.clear();
|
||||
self.file_matches = false;
|
||||
self.path = path;
|
||||
if let Some(p) = self.path.clone() {
|
||||
self.load(&p);
|
||||
}
|
||||
}
|
||||
|
||||
fn header() -> Vec<u8> {
|
||||
let mut h = Vec::new();
|
||||
h.extend_from_slice(SCAN_MAGIC);
|
||||
h.push(SCAN_VERSION);
|
||||
h.extend_from_slice(&(SCAN_TAG.len() as u16).to_le_bytes());
|
||||
h.extend_from_slice(SCAN_TAG.as_bytes());
|
||||
h
|
||||
}
|
||||
|
||||
fn load(&mut self, path: &Path) {
|
||||
let Ok(meta) = std::fs::metadata(path) else { return };
|
||||
if meta.len() > SCAN_MAX_FILE {
|
||||
return;
|
||||
}
|
||||
let Ok(bytes) = std::fs::read(path) else { return };
|
||||
let head = Self::header();
|
||||
let Some(rest) = bytes.strip_prefix(head.as_slice()) else { return };
|
||||
self.file_matches = true;
|
||||
// a half-written last id (a crash while appending) is ignored
|
||||
self.set.extend(rest.as_chunks::<8>().0.iter().map(|c| u64::from_le_bytes(*c)));
|
||||
}
|
||||
|
||||
/// Write what was added since the last save: appended when the file is this detector's, else the whole record
|
||||
/// is written fresh (atomically). A failure is returned and the ids stay pending.
|
||||
pub fn save(&mut self) -> std::io::Result<()> {
|
||||
let Some(path) = self.path.clone() else {
|
||||
self.unsaved.clear();
|
||||
return Ok(());
|
||||
};
|
||||
if self.unsaved.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
if self.file_matches {
|
||||
let mut f = std::fs::OpenOptions::new().append(true).open(&path)?;
|
||||
let buf: Vec<u8> = self.unsaved.iter().flat_map(|id| id.to_le_bytes()).collect();
|
||||
f.write_all(&buf)?;
|
||||
} else {
|
||||
let mut buf = Self::header();
|
||||
buf.extend(self.set.iter().flat_map(|id| id.to_le_bytes()));
|
||||
let part = path.with_extension("part");
|
||||
std::fs::write(&part, &buf)?;
|
||||
std::fs::rename(&part, &path)?;
|
||||
self.file_matches = true;
|
||||
}
|
||||
self.unsaved.clear();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Scanned {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.save();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn rect(x: f64) -> Rect {
|
||||
Rect { x0: x, y0: 0.1, x1: x + 0.2, y1: 0.4 }
|
||||
}
|
||||
|
||||
fn temp(name: &str) -> PathBuf {
|
||||
std::env::temp_dir().join(format!("lc-faceindex-{name}-{}.bin", std::process::id()))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embeddings_survive_a_save_and_load_and_append() {
|
||||
let path = temp("roundtrip");
|
||||
let _ = std::fs::remove_file(&path);
|
||||
let mut a = Index::default();
|
||||
a.reset("m@abc", 3, Some(path.clone()));
|
||||
a.insert(7, &rect(0.1), vec![1.0, 0.0, 0.5]);
|
||||
a.insert(8, &rect(0.3), vec![0.0, 1.0, 0.25]);
|
||||
a.save().unwrap();
|
||||
a.insert(9, &rect(0.5), vec![0.5, 0.5, 0.5]);
|
||||
a.save().unwrap(); // appended
|
||||
let mut b = Index::default();
|
||||
b.reset("m@abc", 3, Some(path.clone()));
|
||||
assert_eq!(b.len(), 3);
|
||||
assert_eq!(b.get(8, &rect(0.3)).map(|e| e.to_vec()), Some(vec![0.0, 1.0, 0.25]));
|
||||
assert!(b.contains(9, &rect(0.5)) && !b.contains(9, &rect(0.6)) && !b.contains(1, &rect(0.5)));
|
||||
let _ = std::fs::remove_file(path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn another_model_a_damaged_file_or_a_torn_tail_are_handled() {
|
||||
let path = temp("other");
|
||||
let _ = std::fs::remove_file(&path);
|
||||
let mut a = Index::default();
|
||||
a.reset("m@abc", 3, Some(path.clone()));
|
||||
a.insert(1, &rect(0.1), vec![1.0, 0.0, 0.0]);
|
||||
a.insert(2, &rect(0.2), vec![0.0, 1.0, 0.0]);
|
||||
a.save().unwrap();
|
||||
// another model: nothing is loaded, and saving replaces the file for the new model
|
||||
let mut other = Index::default();
|
||||
other.reset("n@def", 4, Some(path.clone()));
|
||||
assert_eq!(other.len(), 0);
|
||||
other.insert(5, &rect(0.1), vec![1.0, 0.0, 0.0, 0.0]);
|
||||
other.save().unwrap();
|
||||
let mut again = Index::default();
|
||||
again.reset("n@def", 4, Some(path.clone()));
|
||||
assert_eq!(again.len(), 1);
|
||||
let mut old = Index::default();
|
||||
old.reset("m@abc", 3, Some(path.clone()));
|
||||
assert_eq!(old.len(), 0, "the old model's records are gone");
|
||||
// a torn last record loses only that record
|
||||
let bytes = std::fs::read(&path).unwrap();
|
||||
std::fs::write(&path, &bytes[..bytes.len() - 5]).unwrap();
|
||||
let mut torn = Index::default();
|
||||
torn.reset("n@def", 4, Some(path.clone()));
|
||||
assert_eq!(torn.len(), 0);
|
||||
// garbage of every kind is ignored
|
||||
for junk in [Vec::new(), b"LCFE".to_vec(), b"LCFE\x01\xff\xff".to_vec(), vec![0xff; 200], b"LCFE\x02\x00\x00".to_vec()] {
|
||||
std::fs::write(&path, &junk).unwrap();
|
||||
let mut j = Index::default();
|
||||
j.reset("n@def", 4, Some(path.clone()));
|
||||
assert_eq!(j.len(), 0);
|
||||
}
|
||||
// and a missing file is fine
|
||||
let _ = std::fs::remove_file(&path);
|
||||
let mut m = Index::default();
|
||||
m.reset("n@def", 4, Some(path.clone()));
|
||||
assert_eq!(m.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_finite_boxes_and_vectors_do_not_break_the_keys() {
|
||||
assert_eq!(region_key(&Rect { x0: f64::NAN, y0: f64::INFINITY, x1: 1e300, y1: -1e300 }), [0, 0, 10_000_000, -10_000_000]);
|
||||
let path = temp("nan");
|
||||
let _ = std::fs::remove_file(&path);
|
||||
let mut a = Index::default();
|
||||
a.reset("m@abc", 2, Some(path.clone()));
|
||||
a.insert(1, &rect(0.1), vec![f32::NAN, 1.0]);
|
||||
a.save().unwrap();
|
||||
let mut b = Index::default();
|
||||
b.reset("m@abc", 2, Some(path.clone()));
|
||||
assert_eq!(b.len(), 0, "a non-finite vector is not loaded back");
|
||||
let _ = std::fs::remove_file(path);
|
||||
}
|
||||
}
|
||||
|
||||
/// What the recognition commands keep between calls: the loaded model and the embeddings made with it.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct FacesState {
|
||||
/// `<model id>@<sha256>` and the loaded model.
|
||||
pub embedder: Option<(String, Arc<lightcraft_faces::runtime::Embedder>)>,
|
||||
pub index: Index,
|
||||
/// When the models folder was last looked at (the per-frame pump trusts a look for a second).
|
||||
pub checked: Option<std::time::Instant>,
|
||||
/// Whether recognition was switched on when last looked up, and when.
|
||||
pub enabled_seen: Option<(std::time::Instant, bool)>,
|
||||
pub worker: Option<crate::faces_worker::Worker>,
|
||||
/// After a failure to load the chosen model, the background pump leaves it alone until then.
|
||||
pub retry_at: Option<std::time::Instant>,
|
||||
/// Photos the worker is busy with.
|
||||
pub in_flight: std::collections::HashSet<lightcraft_catalog::PhotoId>,
|
||||
/// Photos still to hand to the worker, and the catalog revision this list was made at.
|
||||
pub queue: Vec<lightcraft_catalog::PhotoId>,
|
||||
pub queue_stamp: Option<u64>,
|
||||
/// When the index was last written to its cache file by the background pump.
|
||||
pub saved_at: Option<std::time::Instant>,
|
||||
/// Bumped when another library is opened: what a worker was making for the old one is dropped.
|
||||
pub epoch: u64,
|
||||
/// Photos already searched for faces.
|
||||
pub scanned: Scanned,
|
||||
/// Photos that could not be searched this session (unreadable files): left alone until the next launch.
|
||||
pub scan_failed: std::collections::HashSet<u64>,
|
||||
}
|
||||
|
||||
impl FacesState {
|
||||
/// Another library is being opened: what was learned about the old one is saved, and nothing of it is kept (photo
|
||||
/// ids are only unique within a library). The loaded model stays; the index is reloaded for the new library.
|
||||
pub fn library_changed(&mut self) {
|
||||
self.index.reset("", 0, None);
|
||||
self.scanned.reset(None);
|
||||
self.scan_failed.clear();
|
||||
self.embedder = None;
|
||||
self.checked = None;
|
||||
self.queue.clear();
|
||||
self.queue_stamp = None;
|
||||
self.in_flight.clear();
|
||||
self.retry_at = None;
|
||||
self.saved_at = None;
|
||||
self.epoch = self.epoch.wrapping_add(1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
//! The heavy part of scanning faces, run off the session thread, on a few threads at once.
|
||||
//!
|
||||
//! For one photo: render it (upright, uncropped, default settings), find its faces with the YuNet detector (when installed)
|
||||
//! so each region can be aligned by landmarks, cut each region out and run the recognition model. A photo that has
|
||||
//! no face regions at all and has not been searched yet also gets its detections back, with their embeddings, so the
|
||||
//! scan decodes each photo once. The session thread only prepares the job (which needs the catalog) and later ingests
|
||||
//! the result, so the window never waits for a model.
|
||||
|
||||
use std::panic::{AssertUnwindSafe, catch_unwind};
|
||||
use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
|
||||
use std::sync::mpsc::{Receiver, Sender, channel};
|
||||
use std::sync::{Arc, Condvar, Mutex, PoisonError};
|
||||
use std::time::Duration;
|
||||
|
||||
use lightcraft_catalog::PhotoId;
|
||||
use lightcraft_faces::align::{Rgb, align_to_template, crop_box};
|
||||
use lightcraft_faces::runtime::Embedder;
|
||||
use lightcraft_faces::yunet::{Detector, Face, Options};
|
||||
use lightcraft_geom::Rect;
|
||||
|
||||
use crate::media::{PreviewLoader, RenderJob};
|
||||
|
||||
/// Long edge of the picture faces are cut from.
|
||||
pub(crate) const EDGE: usize = 2048;
|
||||
/// A box from an XMP sidecar is cut with this much room around it (there are no landmarks to align by).
|
||||
pub(crate) const BOX_PAD: f32 = 1.5;
|
||||
/// How much a detected face must overlap a region to lend it its landmarks.
|
||||
pub(crate) const MATCH_IOU: f32 = 0.3;
|
||||
/// A detection must be this sure to become a face region (the bar `faces.detect` uses).
|
||||
pub(crate) const REGION_SCORE: f32 = 0.6;
|
||||
/// A raw's embedded preview is used for the scan only when its long edge is at least this (a thumbnail is too small to
|
||||
/// find faces in; the raw itself is decoded then).
|
||||
const MIN_PREVIEW_EDGE: usize = 1000;
|
||||
|
||||
/// Most photos ever worked on at once.
|
||||
const MAX_WORKERS: usize = 24;
|
||||
/// What one photo in progress holds: the file, its decoded picture and the copies the detector and recogniser read.
|
||||
/// Measured: peak memory grew by about 180 MB for each extra worker.
|
||||
const WORKER_BYTES: usize = 180 << 20;
|
||||
/// Photos queued behind the ones running, so a worker that finishes never waits for the next frame to be fed.
|
||||
pub(crate) const PREFETCH: usize = 2;
|
||||
|
||||
/// How hard the background scan may work right now, as the app judges from what the user is doing.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum Pace {
|
||||
/// Start nothing new: the user is dragging, typing or scrolling.
|
||||
Pause = 0,
|
||||
/// One photo at a time: the user is around, or the window is minimized.
|
||||
Light = 1,
|
||||
/// Half the machine: the user is idle.
|
||||
Normal = 2,
|
||||
/// Most of the machine: the user is watching the progress.
|
||||
Full = 3,
|
||||
}
|
||||
|
||||
impl Pace {
|
||||
fn from_u8(n: u8) -> Pace {
|
||||
match n {
|
||||
0 => Pace::Pause,
|
||||
1 => Pace::Light,
|
||||
3 => Pace::Full,
|
||||
_ => Pace::Normal,
|
||||
}
|
||||
}
|
||||
|
||||
/// `pause`, `light`, `normal` or `full`; anything else (or nothing) is `normal`.
|
||||
pub fn parse(name: Option<&str>) -> Pace {
|
||||
match name {
|
||||
Some("pause") => Pace::Pause,
|
||||
Some("light") => Pace::Light,
|
||||
Some("full") => Pace::Full,
|
||||
_ => Pace::Normal,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// How many photos are worked on at once at `pace` on a machine with `cores` threads (as the system reports them, which
|
||||
/// respects container limits and affinity). `cap` is a limit the user set (`LIGHTCRAFT_FACE_THREADS`). Memory is not
|
||||
/// decided here: decodes wait at the process-wide memory gate ([`crate::memory::work_gate`]), sized from the RAM.
|
||||
pub(crate) fn workers_for(cores: usize, pace: Pace, cap: Option<usize>) -> usize {
|
||||
let n = match pace {
|
||||
Pace::Pause => return 0,
|
||||
Pace::Light => 1,
|
||||
Pace::Normal => cores / 2,
|
||||
Pace::Full => cores * 4 / 5,
|
||||
}
|
||||
.clamp(1, MAX_WORKERS);
|
||||
cap.map_or(n, |c| n.min(c.clamp(1, MAX_WORKERS)))
|
||||
}
|
||||
|
||||
fn cores() -> usize {
|
||||
std::thread::available_parallelism().map_or(1, |n| n.get())
|
||||
}
|
||||
|
||||
fn cap_from_env() -> Option<usize> {
|
||||
std::env::var("LIGHTCRAFT_FACE_THREADS").ok().and_then(|v| v.trim().parse::<usize>().ok())
|
||||
}
|
||||
|
||||
/// The most photos the memory budget can hold in progress at once: half of it (the rest is the window's, the caches',
|
||||
/// the renders'), a worker at a time. Raising the budget (`LIGHTCRAFT_MEMORY_MB`) raises this.
|
||||
fn memory_workers() -> usize {
|
||||
(crate::memory::budget() / 2 / WORKER_BYTES).max(1)
|
||||
}
|
||||
|
||||
/// Photos to work on at once at `pace` on this machine: by the processor's threads and the pace, and never more than
|
||||
/// the memory allows. A limit the user sets (`LIGHTCRAFT_FACE_THREADS`) replaces the memory one.
|
||||
pub(crate) fn target_workers(pace: Pace) -> usize {
|
||||
workers_for(cores(), pace, Some(cap_from_env().unwrap_or_else(memory_workers)))
|
||||
}
|
||||
|
||||
/// Threads the worker pool starts with: enough for the fullest pace (the idle ones cost nothing).
|
||||
pub(crate) fn max_workers() -> usize {
|
||||
target_workers(Pace::Full)
|
||||
}
|
||||
|
||||
/// Threads for the scan's own parallel work (decoding a picture, developing it) at `pace` on a machine with `cores` threads:
|
||||
/// two when light, half the machine when normal, four fifths when watched. A photo's parallel parts are most of its cost,
|
||||
/// so this, more than the number of photos at once, is how much of the processor the scan takes. Measured on a 32-thread
|
||||
/// desktop (4 photos at once): 8 threads scanned 2.2 photos a second, 16 threads 5.8, 25 threads 6.4.
|
||||
fn pool_threads_for(cores: usize, pace: Pace, cap: Option<usize>) -> usize {
|
||||
let n = match pace {
|
||||
Pace::Pause | Pace::Light => 2,
|
||||
Pace::Normal => cores / 2,
|
||||
Pace::Full => cores * 4 / 5,
|
||||
}
|
||||
.clamp(2, 32);
|
||||
cap.map_or(n, |c| n.min(c.max(1)))
|
||||
}
|
||||
|
||||
/// One photo's faces to embed (and, if `detect`, to find first).
|
||||
pub(crate) struct Prepared {
|
||||
/// `<model id>@<sha256>` the embeddings are for.
|
||||
pub tag: String,
|
||||
/// The session's face epoch when this was prepared: another library has been opened since if it differs.
|
||||
pub epoch: u64,
|
||||
pub id: PhotoId,
|
||||
pub job: RenderJob,
|
||||
/// A raw's embedded preview, tried before the render job (much faster than decoding the raw).
|
||||
pub preview: Option<(String, PreviewLoader)>,
|
||||
pub todo: Vec<Rect>,
|
||||
/// The photo has no face regions and has not been searched: look for faces too.
|
||||
pub detect: bool,
|
||||
/// The YuNet detector, when the user has installed it: faces are aligned by its landmarks, and it finds the faces of
|
||||
/// photos that have none. Without it a scan only embeds the faces the photos already have, cut by their boxes.
|
||||
pub detector: Option<Arc<Detector>>,
|
||||
}
|
||||
|
||||
/// The embeddings made for a [`Prepared`]; `None` for a face that could not be embedded.
|
||||
pub(crate) struct Done {
|
||||
pub tag: String,
|
||||
pub epoch: u64,
|
||||
pub id: PhotoId,
|
||||
/// Each face's region, its embedding (`None` when it could not be made) and the box to show it by.
|
||||
pub results: Vec<(Rect, Option<Vec<f32>>, Rect)>,
|
||||
/// Faces were to be looked for in this photo.
|
||||
pub detect_wanted: bool,
|
||||
/// What the detector found (their embeddings are in `results`); `None` when the photo could not be searched.
|
||||
pub found: Option<Vec<Rect>>,
|
||||
}
|
||||
|
||||
fn overlap(a: &Rect, f: &Face) -> f32 {
|
||||
let (w, h) = ((a.x1.min(f64::from(f.x1)) - a.x0.max(f64::from(f.x0))).max(0.0), (a.y1.min(f64::from(f.y1)) - a.y0.max(f64::from(f.y0))).max(0.0));
|
||||
let inter = w * h;
|
||||
let union = (a.x1 - a.x0) * (a.y1 - a.y0) + f64::from((f.x1 - f.x0) * (f.y1 - f.y0)) - inter;
|
||||
if union > 0.0 { (inter / union) as f32 } else { 0.0 }
|
||||
}
|
||||
|
||||
/// A detection as a region box: inside the photo and not a sliver; `None` for anything else (it is only ever data).
|
||||
fn face_rect(f: &Face) -> Option<Rect> {
|
||||
if ![f.x0, f.y0, f.x1, f.y1].iter().all(|v| v.is_finite()) {
|
||||
return None;
|
||||
}
|
||||
let c = |v: f32| f64::from(v.clamp(0.0, 1.0));
|
||||
let r = Rect { x0: c(f.x0), y0: c(f.y0), x1: c(f.x1), y1: c(f.y1) };
|
||||
(r.x1 - r.x0 > 0.005 && r.y1 - r.y0 > 0.005).then_some(r)
|
||||
}
|
||||
|
||||
/// Do the work for one photo. Never panics: whatever goes wrong, its faces come back as `None`.
|
||||
pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done {
|
||||
let Prepared { tag, epoch, id, job, preview, todo, detect, detector } = p;
|
||||
let fallback = |todo: Vec<Rect>| todo.into_iter().map(|r| (r, None, r)).collect::<Vec<_>>();
|
||||
let rects = todo.clone();
|
||||
let run = catch_unwind(AssertUnwindSafe(|| {
|
||||
let from_preview = preview.and_then(|(path, load)| load(&path, EDGE)).filter(|img| img.width.max(img.height) >= MIN_PREVIEW_EDGE);
|
||||
let image = match from_preview {
|
||||
Some(img) => img,
|
||||
None => match job.run().rendered {
|
||||
Ok(rendered) => rendered.image,
|
||||
Err(_) => return (fallback(todo), None),
|
||||
},
|
||||
};
|
||||
let (w, h) = (image.width, image.height);
|
||||
let rgb: Vec<u8> = image.data.iter().flat_map(|px| [px[0], px[1], px[2]]).collect();
|
||||
// the detector's landmarks, where it finds the same face: faces are aligned by them, not cut by a box
|
||||
let detected = detector.as_deref().and_then(|d| d.detect(&rgb, w, h, &Options { score: 0.5, ..Options::default() }).ok());
|
||||
let found = detected.clone().unwrap_or_default();
|
||||
let img = Rgb { data: &rgb, width: w, height: h };
|
||||
let (iw, ih) = embedder.input_size();
|
||||
let scaled = |f: &Face| f.landmarks.map(|(x, y)| (x * w as f32, y * h as f32));
|
||||
let mut results: Vec<(Rect, Option<Vec<f32>>, Rect)> = todo
|
||||
.into_iter()
|
||||
.map(|rect| {
|
||||
let best = found.iter().map(|f| (overlap(&rect, f), f)).filter(|(o, _)| *o >= MATCH_IOU).max_by(|a, b| a.0.total_cmp(&b.0));
|
||||
let crop = match best {
|
||||
Some((_, f)) => align_to_template(&img, &scaled(f), iw, ih),
|
||||
None => crop_box(
|
||||
&img,
|
||||
rect.x0 as f32 * w as f32,
|
||||
rect.y0 as f32 * h as f32,
|
||||
rect.x1 as f32 * w as f32,
|
||||
rect.y1 as f32 * h as f32,
|
||||
BOX_PAD,
|
||||
iw,
|
||||
ih,
|
||||
),
|
||||
};
|
||||
let v = crop.and_then(|c| embedder.embed(&c).ok());
|
||||
// shown by the detector's box where it found this face: every face then has the same tightness
|
||||
let view = best.and_then(|(_, f)| face_rect(f)).unwrap_or(rect);
|
||||
(rect, v, view)
|
||||
})
|
||||
.collect();
|
||||
// a photo with no face regions: what the detector is sure of becomes its faces, embedded in the same pass
|
||||
let new_faces = if detect {
|
||||
detected.map(|faces| {
|
||||
let mut rects = Vec::new();
|
||||
for f in faces.iter().filter(|f| f.score >= REGION_SCORE) {
|
||||
let Some(rect) = face_rect(f) else { continue };
|
||||
let v = align_to_template(&img, &scaled(f), iw, ih).and_then(|c| embedder.embed(&c).ok());
|
||||
results.push((rect, v, rect));
|
||||
rects.push(rect);
|
||||
}
|
||||
rects
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
(results, new_faces)
|
||||
}));
|
||||
let (results, found) = run.unwrap_or_else(|_| (fallback(rects), None));
|
||||
Done { tag, epoch, id, results, detect_wanted: detect, found }
|
||||
}
|
||||
|
||||
type Job = (Arc<Embedder>, Prepared);
|
||||
|
||||
/// How many photos may be worked on at once, changed at any moment by the pace: a thread that has a job waits here until
|
||||
/// fewer than that are running. A pace of zero holds every job where it is, with nothing decoded.
|
||||
struct Limit {
|
||||
allowed: AtomicUsize,
|
||||
/// The pace the limit was last set for (a [`Pace`] as a number): the pool a photo's parallel work uses.
|
||||
pace: AtomicU8,
|
||||
running: Mutex<usize>,
|
||||
wake: Condvar,
|
||||
}
|
||||
|
||||
impl Limit {
|
||||
fn new() -> Limit {
|
||||
Limit { allowed: AtomicUsize::new(0), pace: AtomicU8::new(Pace::Normal as u8), running: Mutex::new(0), wake: Condvar::new() }
|
||||
}
|
||||
|
||||
fn set(&self, n: usize, pace: Pace) {
|
||||
self.pace.store(pace as u8, Ordering::Relaxed);
|
||||
self.allowed.store(n, Ordering::Relaxed);
|
||||
self.wake.notify_all();
|
||||
}
|
||||
|
||||
fn pace(&self) -> Pace {
|
||||
Pace::from_u8(self.pace.load(Ordering::Relaxed))
|
||||
}
|
||||
|
||||
/// Wait for a place; the limit is looked at again whenever it changes and at least every 100 ms.
|
||||
fn enter(&self) {
|
||||
let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
while *running >= self.allowed.load(Ordering::Relaxed) {
|
||||
running = self.wake.wait_timeout(running, Duration::from_millis(100)).unwrap_or_else(PoisonError::into_inner).0;
|
||||
}
|
||||
*running += 1;
|
||||
}
|
||||
|
||||
fn leave(&self) {
|
||||
let mut running = self.running.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
*running = running.saturating_sub(1);
|
||||
drop(running);
|
||||
self.wake.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
/// The threads a photo's parallel work (decoding, developing a picture) runs on at `pace`: one pool for each of three sizes
|
||||
/// ([`pool_threads_for`]), made when first needed, whose idle threads cost nothing. rayon's global pool serves the
|
||||
/// interactive work (the loupe, exports) and has no priorities, so scan work queued there would make a slider drag wait
|
||||
/// behind it; here a light scan can use only two threads however large the machine is.
|
||||
fn scan_pool(pace: Pace) -> Option<&'static rayon::ThreadPool> {
|
||||
static LIGHT: std::sync::OnceLock<Option<rayon::ThreadPool>> = std::sync::OnceLock::new();
|
||||
static NORMAL: std::sync::OnceLock<Option<rayon::ThreadPool>> = std::sync::OnceLock::new();
|
||||
static FULL: std::sync::OnceLock<Option<rayon::ThreadPool>> = std::sync::OnceLock::new();
|
||||
let (cell, pace) = match pace {
|
||||
Pace::Pause | Pace::Light => (&LIGHT, Pace::Light),
|
||||
Pace::Normal => (&NORMAL, Pace::Normal),
|
||||
Pace::Full => (&FULL, Pace::Full),
|
||||
};
|
||||
cell.get_or_init(|| {
|
||||
rayon::ThreadPoolBuilder::new()
|
||||
.num_threads(pool_threads_for(cores(), pace, cap_from_env()))
|
||||
.thread_name(move |i| format!("lightcraft-faces-{}-{i}", pace as u8))
|
||||
.build()
|
||||
.ok()
|
||||
})
|
||||
.as_ref()
|
||||
}
|
||||
|
||||
/// Run one photo on the pool for the pace, so its parallel parts stay there.
|
||||
///
|
||||
/// It must NOT wait at the memory gate ([`crate::memory::work_gate`], which background decodes do): a job running on a
|
||||
/// pool thread that waits inside a nested rayon call can pick up another job to run on top of itself, and if that one waits
|
||||
/// for memory the first holds, neither can finish (seen as a few photos that never come back). The memory a scan may use is
|
||||
/// instead limited by how many photos run at once ([`memory_workers`]); its decodes count as interactive ones at the gate,
|
||||
/// which never wait.
|
||||
fn process_in_background(prepared: Prepared, embedder: &Embedder, pace: Pace) -> Done {
|
||||
let work = || process(prepared, embedder);
|
||||
match scan_pool(pace) {
|
||||
Some(pool) => pool.install(work),
|
||||
None => work(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Threads that run [`process`] for the jobs they are given, each on one photo at a time.
|
||||
pub(crate) struct Worker {
|
||||
jobs: Sender<Job>,
|
||||
done: Receiver<Done>,
|
||||
limit: Arc<Limit>,
|
||||
}
|
||||
|
||||
impl Worker {
|
||||
pub fn start(threads: usize) -> Option<Worker> {
|
||||
let (jobs, job_rx) = channel::<Job>();
|
||||
let (done_tx, done) = channel::<Done>();
|
||||
let job_rx = Arc::new(Mutex::new(job_rx));
|
||||
let limit = Arc::new(Limit::new());
|
||||
let mut started = 0;
|
||||
for i in 0..threads.clamp(1, MAX_WORKERS) {
|
||||
let (rx, tx, limit) = (job_rx.clone(), done_tx.clone(), limit.clone());
|
||||
let spawned = std::thread::Builder::new().name(format!("lightcraft-faces-{i}")).spawn(move || {
|
||||
loop {
|
||||
// the lock is held only while waiting for a job, never while working on one; the threads end when
|
||||
// the session (the only sender) is dropped
|
||||
let next = rx.lock().unwrap_or_else(PoisonError::into_inner).recv();
|
||||
let Ok((embedder, prepared)) = next else { break };
|
||||
limit.enter();
|
||||
let done = process_in_background(prepared, &embedder, limit.pace());
|
||||
limit.leave();
|
||||
if tx.send(done).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
if spawned.is_ok() {
|
||||
started += 1;
|
||||
}
|
||||
}
|
||||
(started > 0).then_some(Worker { jobs, done, limit })
|
||||
}
|
||||
|
||||
/// How many photos may be worked on at once from now on (what is running is not interrupted).
|
||||
pub fn allow(&self, n: usize, pace: Pace) {
|
||||
self.limit.set(n, pace);
|
||||
}
|
||||
|
||||
pub fn submit(&self, embedder: Arc<Embedder>, prepared: Prepared) -> bool {
|
||||
self.jobs.send((embedder, prepared)).is_ok()
|
||||
}
|
||||
|
||||
/// Everything finished since the last call.
|
||||
pub fn finished(&self) -> Vec<Done> {
|
||||
self.done.try_iter().collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn the_pace_sets_how_much_of_the_machine_the_scan_may_use() {
|
||||
use Pace::*;
|
||||
// a 32-thread desktop: nothing while the user works, one photo around them, half when idle, most when watched
|
||||
assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(32, p, None)), [0, 1, 16, 24]);
|
||||
assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(8, p, None)), [0, 1, 4, 6]);
|
||||
assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(4, p, None)), [0, 1, 2, 3]);
|
||||
// small machines still get one worker (except while paused)
|
||||
assert_eq!([Pause, Light, Normal, Full].map(|p| workers_for(1, p, None)), [0, 1, 1, 1]);
|
||||
assert_eq!(workers_for(0, Full, None), 1);
|
||||
assert_eq!(workers_for(10_000, Full, None), MAX_WORKERS);
|
||||
// a limit set by the user caps every pace, and is itself kept sane
|
||||
assert_eq!([Light, Normal, Full].map(|p| workers_for(32, p, Some(3))), [1, 3, 3]);
|
||||
assert_eq!((workers_for(32, Full, Some(0)), workers_for(32, Pause, Some(5))), (1, 0));
|
||||
assert_eq!(workers_for(32, Full, Some(usize::MAX)), 24);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_limit_holds_threads_to_the_pace_and_can_change_at_any_moment() {
|
||||
let limit = Arc::new(Limit::new());
|
||||
let (now, peak) = (Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
|
||||
let started = Arc::new(AtomicUsize::new(0));
|
||||
let threads: Vec<_> = (0..6)
|
||||
.map(|_| {
|
||||
let (limit, now, peak, started) = (limit.clone(), now.clone(), peak.clone(), started.clone());
|
||||
std::thread::spawn(move || {
|
||||
limit.enter();
|
||||
started.fetch_add(1, Ordering::SeqCst);
|
||||
let n = now.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
peak.fetch_max(n, Ordering::SeqCst);
|
||||
std::thread::sleep(Duration::from_millis(40));
|
||||
now.fetch_sub(1, Ordering::SeqCst);
|
||||
limit.leave();
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
// paused: nothing starts
|
||||
std::thread::sleep(Duration::from_millis(150));
|
||||
assert_eq!(started.load(Ordering::SeqCst), 0);
|
||||
// two at a time, never more
|
||||
limit.set(2, Pace::Normal);
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
assert!(started.load(Ordering::SeqCst) >= 2);
|
||||
// raised, the rest go; every thread finishes
|
||||
limit.set(6, Pace::Full);
|
||||
for t in threads {
|
||||
t.join().unwrap();
|
||||
}
|
||||
assert_eq!(started.load(Ordering::SeqCst), 6);
|
||||
assert!(peak.load(Ordering::SeqCst) >= 2, "it did run in parallel");
|
||||
assert_eq!(*limit.running.lock().unwrap(), 0, "every place was given back");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_limits_the_workers_and_a_users_limit_replaces_it() {
|
||||
// the default budget (a quarter of the RAM, at most 1.5 GiB) holds a few photos in progress, not dozens
|
||||
assert!((1..=8).contains(&memory_workers()), "{}", memory_workers());
|
||||
assert!(target_workers(Pace::Full) >= 1);
|
||||
assert_eq!(target_workers(Pace::Pause), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pace_names_are_forgiving() {
|
||||
assert_eq!(["pause", "light", "normal", "full"].map(|n| Pace::parse(Some(n))), [Pace::Pause, Pace::Light, Pace::Normal, Pace::Full]);
|
||||
assert_eq!(
|
||||
(Pace::parse(None), Pace::parse(Some("")), Pace::parse(Some("FULL")), Pace::parse(Some("turbo"))),
|
||||
(Pace::Normal, Pace::Normal, Pace::Normal, Pace::Normal)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_scans_parallel_work_gets_a_pool_the_size_of_the_pace() {
|
||||
use Pace::*;
|
||||
assert_eq!([Pause, Light, Normal, Full].map(|p| pool_threads_for(32, p, None)), [2, 2, 16, 25]);
|
||||
assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(8, p, None)), [2, 4, 6]);
|
||||
assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(2, p, None)), [2, 2, 2]);
|
||||
assert_eq!(pool_threads_for(10_000, Full, None), 32);
|
||||
assert_eq!([Light, Normal, Full].map(|p| pool_threads_for(32, p, Some(3))), [2, 3, 3]);
|
||||
// the pools exist and are the sizes asked for
|
||||
for pace in [Light, Normal, Full] {
|
||||
let pool = scan_pool(pace).expect("a pool can be built");
|
||||
assert_eq!(pool.install(rayon::current_num_threads), pool_threads_for(cores(), pace, cap_from_env()));
|
||||
}
|
||||
// a paused scan has the light pool; the pace is remembered as a number and read back
|
||||
assert_eq!(scan_pool(Pause).map(|p| p as *const _), scan_pool(Light).map(|p| p as *const _));
|
||||
assert_eq!([0, 1, 2, 3, 9].map(Pace::from_u8), [Pause, Light, Normal, Full, Normal]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detections_become_boxes_only_when_they_make_sense() {
|
||||
let f = |x0, y0, x1, y1| Face { x0, y0, x1, y1, score: 0.9, landmarks: [(0.0, 0.0); 5] };
|
||||
assert!(face_rect(&f(0.2, 0.2, 0.4, 0.5)).is_some());
|
||||
// a box that spills over the edge is brought back inside
|
||||
let r = face_rect(&f(-0.1, 0.1, 0.3, 1.2)).unwrap();
|
||||
assert_eq!((r.x0, r.y1), (0.0, 1.0));
|
||||
// slivers, inverted boxes and non-numbers are dropped, never turned into regions
|
||||
for bad in [f(0.2, 0.2, 0.2, 0.5), f(0.4, 0.2, 0.2, 0.5), f(f32::NAN, 0.2, 0.4, 0.5), f(0.2, 0.2, f32::INFINITY, 0.5), f(1.5, 0.2, 2.0, 0.5)]
|
||||
{
|
||||
assert!(face_rect(&bad).is_none(), "{bad:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
+302
-39
@@ -11,7 +11,7 @@ use lightcraft_pipeline::SourceInfo;
|
||||
use lightcraft_raster::Rgb32f;
|
||||
use lightcraft_raster::resample::{Filter, fit};
|
||||
|
||||
use crate::media::{FileLoader, FileProbe, PreviewLoader, ProbeInfo};
|
||||
use crate::media::{DenoiseSpec, FileLoader, FileProbe, PairLoader, PreviewLoader, ProbeInfo};
|
||||
|
||||
fn meta_of(m: &lightcraft_meta::Metadata) -> (Meta, Option<String>) {
|
||||
let shutter = m.exposure_time.map(|t| if t >= 1.0 { format!("{t:.0}") } else { format!("1/{:.0}", 1.0 / t) }).unwrap_or_default();
|
||||
@@ -220,22 +220,56 @@ pub fn bin_factor(raw: &lightcraft_raw::RawImage, max_edge: usize) -> Option<usi
|
||||
/// instead of each one waking the pool from outside and waiting for it (which costs more than the
|
||||
/// loops themselves when the machine is busy).
|
||||
pub fn load_bytes(bytes: &[u8], max_edge: usize) -> Result<(Rgb32f, SourceInfo), String> {
|
||||
rayon::scope(|_| load_bytes_now(std::borrow::Cow::Borrowed(bytes), max_edge))
|
||||
rayon::scope(|_| load_bytes_now(std::borrow::Cow::Borrowed(bytes), max_edge, None).map(|(img, _, info)| (img, info)))
|
||||
}
|
||||
|
||||
/// [`load_bytes`] taking the file's bytes: a raw file's bytes are freed as soon as it is decoded
|
||||
/// (less memory held while it is developed).
|
||||
pub fn load_vec(bytes: Vec<u8>, max_edge: usize) -> Result<(Rgb32f, SourceInfo), String> {
|
||||
rayon::scope(move |_| load_bytes_now(std::borrow::Cow::Owned(bytes), max_edge))
|
||||
rayon::scope(move |_| load_bytes_now(std::borrow::Cow::Owned(bytes), max_edge, None).map(|(img, _, info)| (img, info)))
|
||||
}
|
||||
|
||||
fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result<(Rgb32f, SourceInfo), String> {
|
||||
/// [`load_vec`] that also develops the photo's cached denoised picture (`spec`) from the same decode: the plain
|
||||
/// picture, the denoised one (`None` when it cannot be read or does not fit this raw: the plain one is then all
|
||||
/// there is) and the source facts. Both pictures are the same size.
|
||||
pub fn load_vec_pair(bytes: Vec<u8>, max_edge: usize, spec: &DenoiseSpec) -> Result<(Rgb32f, Option<Rgb32f>, SourceInfo), String> {
|
||||
rayon::scope(move |_| load_bytes_now(std::borrow::Cow::Owned(bytes), max_edge, Some(spec)))
|
||||
}
|
||||
|
||||
/// The denoised camera RGB of `raw` from its cached product, as the plain development would have it: the default
|
||||
/// crop, and `factor` × `factor` blocks averaged into one pixel when the plain one was binned. `None` when the product
|
||||
/// is missing, damaged, made from something else or not the size of this raw's active area.
|
||||
fn denoised_camera_rgb(raw: &lightcraft_raw::RawImage, spec: &DenoiseSpec, factor: usize) -> Option<Rgb32f> {
|
||||
use lightcraft_denoise::product::{self, Window};
|
||||
let a = raw.active_area;
|
||||
let header = product::read_header_at(&spec.product).ok()?;
|
||||
if header.key != spec.key || (header.width, header.height) != (a.width, a.height) {
|
||||
return None;
|
||||
}
|
||||
if raw.opcodes.list3.is_empty() {
|
||||
let c = raw.develop_crop(a.width, a.height);
|
||||
let window = Window { x: c.x, y: c.y, width: c.width, height: c.height };
|
||||
product::read_window(&spec.product, &spec.key, Some(window), factor, Some(HIGHLIGHT_CLIP)).ok()
|
||||
} else {
|
||||
// (the plain development is not binned either when the file has opcodes of this kind: `factor` is 1)
|
||||
let whole = product::read(&spec.product, &spec.key, 1).ok()?;
|
||||
Some(raw.finish_demosaiced(whole))
|
||||
}
|
||||
}
|
||||
|
||||
fn load_bytes_now(
|
||||
bytes: std::borrow::Cow<'_, [u8]>,
|
||||
max_edge: usize,
|
||||
denoise: Option<&DenoiseSpec>,
|
||||
) -> Result<(Rgb32f, Option<Rgb32f>, SourceInfo), String> {
|
||||
if lightcraft_raw::probe(&bytes).is_some() {
|
||||
let mut raw = match lightcraft_raw::decode(&bytes).map_err(|e| preview_reason(&bytes, e)) {
|
||||
Ok(r) => r,
|
||||
Err(Ok(why)) => {
|
||||
// show the camera's embedded JPEG (rendered, not raw) until the variant is supported
|
||||
return load_embedded_preview(&bytes, max_edge).ok_or(format!("unsupported raw ({why}) without an embedded preview"));
|
||||
return load_embedded_preview(&bytes, max_edge)
|
||||
.map(|(img, info)| (img, None, info))
|
||||
.ok_or(format!("unsupported raw ({why}) without an embedded preview"));
|
||||
}
|
||||
Err(Err(e)) => return Err(e),
|
||||
};
|
||||
@@ -245,53 +279,63 @@ fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result<
|
||||
raw.opcodes.list3.retain(|op| !op.is_lens_correction());
|
||||
let xy = lightcraft_raw::color::as_shot_white_xy(&raw);
|
||||
let t = lightcraft_raw::color::camera_transform(&raw, xy);
|
||||
let camera_look = crate::camera_preview::fit_preview(&raw, &bytes, &t);
|
||||
// the source's segmentation mattes (DNG semantic masks), read while the preview is fitted
|
||||
let (camera_look, mattes) = rayon::join(|| crate::camera_preview::fit_preview(&raw, &bytes, &t), || dng_mattes(&bytes, &raw));
|
||||
drop(bytes);
|
||||
// Previews and thumbnails bin the mosaic straight to (about) the size they need; only
|
||||
// larger levels (exports, 1:1) demosaic the whole sensor.
|
||||
let t0 = web_time::Instant::now();
|
||||
let binned = match bin_factor(&raw, max_edge) {
|
||||
let bin = bin_factor(&raw, max_edge);
|
||||
let binned = match bin {
|
||||
Some(k) => raw.develop_binned(k, HIGHLIGHT_CLIP).map_err(|e| e.to_string())?,
|
||||
None => None,
|
||||
};
|
||||
let mut img = match binned {
|
||||
let factor = if binned.is_some() { bin.unwrap_or(1) } else { 1 };
|
||||
let img = match binned {
|
||||
Some(img) => img,
|
||||
None => {
|
||||
let method = if max_edge <= 600 { lightcraft_raw::Method::Bilinear } else { lightcraft_raw::Method::Ahd };
|
||||
raw.develop(method).map_err(|e| e.to_string())?
|
||||
}
|
||||
};
|
||||
// the same development from the denoised mosaic, when the photo has been denoised
|
||||
let twin = denoise.and_then(|spec| denoised_camera_rgb(&raw, spec, factor)).filter(|d| d.width == img.width && d.height == img.height);
|
||||
// the samples aren't needed any more (the colour model below reads only the tags)
|
||||
raw.data = lightcraft_raw::RawData::U16(Vec::new());
|
||||
let mut stages = vec![("develop", t0.elapsed())];
|
||||
stages.push(("transform", t0.elapsed()));
|
||||
lightcraft_raw::highlight::reconstruct(&mut img, t.wb, HIGHLIGHT_CLIP);
|
||||
stages.push(("highlights", t0.elapsed()));
|
||||
let wb = t.wb;
|
||||
let m = camera_look.as_ref().map(|p| p.matrix.mul(&t.matrix)).unwrap_or(t.matrix).to_f32();
|
||||
let hue_sat = camera_look.as_ref().and_then(|p| p.hue_sat.as_ref()).and_then(crate::camera_preview::HueSat::new);
|
||||
let gain = 2f32.powf(t.baseline_exposure as f32);
|
||||
let wb = t.wb;
|
||||
// A DNG's own profile look (hue/saturation map, look table), DNG spec chapter 6.
|
||||
let tables = lightcraft_raw::profile::ProfileTables::new(&raw.color.profile, lightcraft_raw::color::illuminant_weight(&raw.color, xy));
|
||||
img.map_in_place(|p| {
|
||||
let c = [p[0] * wb[0], p[1] * wb[1], p[2] * wb[2]];
|
||||
let rgb = [
|
||||
m[0][0] * c[0] + m[0][1] * c[1] + m[0][2] * c[2],
|
||||
m[1][0] * c[0] + m[1][1] * c[1] + m[1][2] * c[2],
|
||||
m[2][0] * c[0] + m[2][1] * c[1] + m[2][2] * c[2],
|
||||
];
|
||||
let rgb = match &tables {
|
||||
Some(tables) => tables.apply(rgb, gain),
|
||||
None => rgb.map(|v| v * gain),
|
||||
};
|
||||
// after the baseline exposure, as when it was fitted
|
||||
hue_sat.as_ref().map_or(rgb, |h| h.apply(rgb)).map(|v| v.max(0.0))
|
||||
});
|
||||
stages.push(("colour", t0.elapsed()));
|
||||
let img = fit(&img, max_edge, max_edge, Filter::Box);
|
||||
stages.push(("fit", t0.elapsed()));
|
||||
let img = img.into_oriented(raw.orientation);
|
||||
stages.push(("orient", t0.elapsed()));
|
||||
// camera RGB → what the pipeline takes: clipped highlights rebuilt, white balance and the colour model applied,
|
||||
// fitted to `max_edge` and upright (the same for the plain and the denoised picture)
|
||||
let finish = |mut img: Rgb32f| {
|
||||
let mut stages = vec![("develop", t0.elapsed())];
|
||||
stages.push(("transform", t0.elapsed()));
|
||||
lightcraft_raw::highlight::reconstruct(&mut img, wb, HIGHLIGHT_CLIP);
|
||||
stages.push(("highlights", t0.elapsed()));
|
||||
img.map_in_place(|p| {
|
||||
let c = [p[0] * wb[0], p[1] * wb[1], p[2] * wb[2]];
|
||||
let rgb = [
|
||||
m[0][0] * c[0] + m[0][1] * c[1] + m[0][2] * c[2],
|
||||
m[1][0] * c[0] + m[1][1] * c[1] + m[1][2] * c[2],
|
||||
m[2][0] * c[0] + m[2][1] * c[1] + m[2][2] * c[2],
|
||||
];
|
||||
let rgb = match &tables {
|
||||
Some(tables) => tables.apply(rgb, gain),
|
||||
None => rgb.map(|v| v * gain),
|
||||
};
|
||||
hue_sat.as_ref().map_or(rgb, |h| h.apply(rgb)).map(|v| v.max(0.0))
|
||||
});
|
||||
stages.push(("colour", t0.elapsed()));
|
||||
let img = fit(&img, max_edge, max_edge, Filter::Box);
|
||||
stages.push(("fit", t0.elapsed()));
|
||||
let img = img.into_oriented(raw.orientation);
|
||||
stages.push(("orient", t0.elapsed()));
|
||||
(img, stages)
|
||||
};
|
||||
let (img, stages) = finish(img);
|
||||
if lightcraft_pipeline::profiling() {
|
||||
let mut prev = std::time::Duration::ZERO;
|
||||
let parts: Vec<String> = stages
|
||||
@@ -304,17 +348,64 @@ fn load_bytes_now(bytes: std::borrow::Cow<'_, [u8]>, max_edge: usize) -> Result<
|
||||
.collect();
|
||||
eprintln!("[profile] raw source {}×{} (max {max_edge}, ms after decode): {}", img.width, img.height, parts.join(", "));
|
||||
}
|
||||
let twin = twin.map(|picture| finish(picture).0);
|
||||
let (temp, tint) = xy_to_temp_tint(xy);
|
||||
let relative = crate::camera_preview::file_local_look(raw.format) && t.matrix_is_fallback;
|
||||
// White balance re-evaluates the file's own colour model (when it has one and no
|
||||
// file-local look matrix sits on top of it)
|
||||
let camera_color = (!t.matrix_is_fallback && camera_look.is_none()).then(|| {
|
||||
let tags = lightcraft_raw::ColorData { profile: Default::default(), ..raw.color.clone() };
|
||||
Arc::new(lightcraft_pipeline::CameraColor { tags, developed_for: xy })
|
||||
});
|
||||
let camera_tone = camera_look.as_ref().map(|p| p.tone).or_else(|| raw.color.profile.tone_curve.as_ref().and_then(dng_tone_curve));
|
||||
let (temp, tint) = if relative { (6500.0, 0.0) } else { (temp.round(), tint.round()) };
|
||||
return Ok((img, SourceInfo { raw: true, as_shot_temp: temp, as_shot_tint: tint, lens, relative_wb: relative, camera_tone }));
|
||||
return Ok((
|
||||
img,
|
||||
twin,
|
||||
SourceInfo { raw: true, as_shot_temp: temp, as_shot_tint: tint, lens, relative_wb: relative, camera_color, camera_tone, mattes },
|
||||
));
|
||||
}
|
||||
let d = lightcraft_codecs::decode(&bytes, fit_box(max_edge)).map_err(|e| e.to_string())?;
|
||||
drop(bytes);
|
||||
let img = d.to_working();
|
||||
let img = if img.width.max(img.height) > max_edge { fit(&img, max_edge, max_edge, Filter::Mitchell) } else { img };
|
||||
Ok((img.into_oriented(Orientation::from_exif(d.orientation)), SourceInfo::default()))
|
||||
Ok((img.into_oriented(Orientation::from_exif(d.orientation)), None, SourceInfo::default()))
|
||||
}
|
||||
|
||||
/// The semantic masks of a DNG that AI masks understand, over the developed image (default crop,
|
||||
/// oriented like it). Person mattes without a confidently selected pixel are left out (the
|
||||
/// Subject mask then falls back to its heuristic); a sky matte counts even when empty (no sky).
|
||||
fn dng_mattes(bytes: &[u8], raw: &lightcraft_raw::RawImage) -> Option<Arc<lightcraft_pipeline::masks::Mattes>> {
|
||||
use lightcraft_pipeline::masks::{MatteKind, Mattes};
|
||||
if raw.format != lightcraft_raw::RawFormat::Dng {
|
||||
return None;
|
||||
}
|
||||
let mut mattes = Mattes::default();
|
||||
for m in lightcraft_raw::semantic_masks(bytes) {
|
||||
let Some(kind) = matte_kind(&m.name) else { continue };
|
||||
let Some(img) = m.developed(raw.active_area, raw.crop) else { continue };
|
||||
if kind != MatteKind::Sky && !img.data.iter().any(|&v| v >= 128) {
|
||||
continue;
|
||||
}
|
||||
mattes.push(kind, img.into_oriented(raw.orientation));
|
||||
}
|
||||
(!mattes.is_empty()).then(|| Arc::new(mattes))
|
||||
}
|
||||
|
||||
/// What a semantic mask selects, by its `SemanticName`: Apple's (iPhone ProRAW) are named after
|
||||
/// their auxiliary image types, `urn:com:apple:photo:<year>:aux:<type>`. Other names: `None`.
|
||||
fn matte_kind(name: &str) -> Option<lightcraft_pipeline::masks::MatteKind> {
|
||||
use lightcraft_pipeline::masks::MatteKind::*;
|
||||
let kind = name.strip_prefix("urn:com:apple:photo:")?.rsplit(':').next()?;
|
||||
Some(match kind {
|
||||
"semanticskymatte" => Sky,
|
||||
"semanticskinmatte" => Skin,
|
||||
"semantichairmatte" => Hair,
|
||||
"semanticteethmatte" => Teeth,
|
||||
"semanticglassesmatte" => Glasses,
|
||||
"portraiteffectsmatte" => Person,
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
/// A DNG `ProfileToneCurve` (linear in, linear out, 1.0 = white after exposure compensation) as the
|
||||
@@ -404,8 +495,31 @@ fn fit_box(max_edge: usize) -> lightcraft_codecs::DecodeOptions {
|
||||
}
|
||||
|
||||
/// Filesystem-backed embedded-preview hook (native).
|
||||
///
|
||||
/// Reads the whole file to find the preview, so it holds the memory gate for twice the file's size (the file and the
|
||||
/// decoded preview): background work (the face scan) waits for room, interactive work is counted and never waits.
|
||||
pub fn fs_preview_loader() -> PreviewLoader {
|
||||
Arc::new(|path: &str, max_edge: usize| embedded_preview_srgb(&std::fs::read(path).ok()?, max_edge))
|
||||
Arc::new(|path: &str, max_edge: usize| {
|
||||
let len = std::fs::metadata(path).map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX)).unwrap_or(0);
|
||||
let gate = crate::memory::work_gate();
|
||||
let _permit = if crate::memory::is_background() { gate.acquire(len.saturating_mul(2)) } else { gate.acquire_urgent(len.saturating_mul(2)) };
|
||||
embedded_preview_srgb(&std::fs::read(path).ok()?, max_edge)
|
||||
})
|
||||
}
|
||||
|
||||
/// The filesystem-backed [`PairLoader`]: the file is read and decoded once, and developed as usual and from the photo's
|
||||
/// denoised picture. It holds two pictures at once, so it counts for more at the memory gate than a plain load.
|
||||
pub fn fs_pair_loader() -> PairLoader {
|
||||
Arc::new(|path: &str, max_edge: usize, spec: &DenoiseSpec| {
|
||||
let len = std::fs::metadata(path).map(|m| m.len() as usize).unwrap_or(0);
|
||||
let weight = len * if max_edge <= crate::media::SourceLevel::Thumb.max_edge() { 4 } else { 10 };
|
||||
let gate = crate::memory::work_gate();
|
||||
let _permit = if crate::memory::is_background() { gate.acquire(weight) } else { gate.acquire_urgent(weight) };
|
||||
let bytes = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?;
|
||||
let r = load_vec_pair(bytes, max_edge, spec);
|
||||
crate::memory::release();
|
||||
r
|
||||
})
|
||||
}
|
||||
|
||||
/// Filesystem-backed hooks (native). On the web the host installs bytes-based hooks instead.
|
||||
@@ -417,7 +531,7 @@ pub fn fs_preview_loader() -> PreviewLoader {
|
||||
/// yields to them.
|
||||
pub fn fs_hooks() -> (FileLoader, FileProbe) {
|
||||
let loader: FileLoader = Arc::new(|path: &str, max_edge: usize| {
|
||||
let len = std::fs::metadata(path).map(|m| m.len() as usize).unwrap_or(0);
|
||||
let len = std::fs::metadata(path).map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX)).unwrap_or(0);
|
||||
let weight = len * if max_edge <= crate::media::SourceLevel::Thumb.max_edge() { 3 } else { 6 };
|
||||
let gate = crate::memory::work_gate();
|
||||
let _permit = if crate::memory::is_background() { gate.acquire(weight) } else { gate.acquire_urgent(weight) };
|
||||
@@ -428,7 +542,7 @@ pub fn fs_hooks() -> (FileLoader, FileProbe) {
|
||||
r
|
||||
});
|
||||
let probe: FileProbe = Arc::new(|path: &str| {
|
||||
let len = std::fs::metadata(path).map(|m| m.len() as usize).unwrap_or(0);
|
||||
let len = std::fs::metadata(path).map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX)).unwrap_or(0);
|
||||
let _permit = crate::memory::work_gate().acquire(len);
|
||||
let bytes = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?;
|
||||
probe_bytes(path, &bytes)
|
||||
@@ -443,6 +557,7 @@ impl crate::Session {
|
||||
self.media.file_loader = Some(l);
|
||||
self.media.file_probe = Some(p);
|
||||
self.media.preview_loader = Some(fs_preview_loader());
|
||||
self.media.denoise.loader = Some(fs_pair_loader());
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -483,6 +598,37 @@ mod tests {
|
||||
assert!(probe_bytes("x.jpg", b"not an image").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_denoise_uses_the_same_white_balance_profile_and_orientation_as_plain_raw() {
|
||||
use lightcraft_denoise::product;
|
||||
use lightcraft_raw::{DngWriteOptions, Method, Orientation, profile::HsvTable, write_dng};
|
||||
let dir = crate::tests_xmp::temp_dir("denoise-color-finish");
|
||||
let fixture = crate::tests_xmp::synthetic_dng_with(None, Default::default());
|
||||
let mut raw = lightcraft_raw::decode(&fixture).unwrap();
|
||||
raw.orientation = Orientation::Rotate90;
|
||||
raw.color.baseline_exposure = 0.5;
|
||||
raw.color.profile.look_table =
|
||||
Some(HsvTable { hue_divisions: 1, sat_divisions: 2, val_divisions: 1, data: vec![[23.0, 0.65, 0.8]; 2], srgb_value: false });
|
||||
let bytes = write_dng(&raw, &DngWriteOptions::default()).unwrap();
|
||||
let decoded = lightcraft_raw::decode(&bytes).unwrap();
|
||||
assert!(decoded.color.profile.look_table.is_some());
|
||||
// A cache containing the ordinary demosaic must undergo the identical downstream finish.
|
||||
let camera_rgb = decoded.develop(Method::Ahd).unwrap();
|
||||
let key = "synthetic-color-finish";
|
||||
let path = dir.join("same.lcdn");
|
||||
product::write(&path, &camera_rgb, key).unwrap();
|
||||
let spec = DenoiseSpec { product: path, key: key.into(), make: None };
|
||||
let (ordinary, _) = load_bytes(&bytes, 2048).unwrap();
|
||||
let (plain, cached, _) = load_vec_pair(bytes, 2048, &spec).unwrap();
|
||||
assert_eq!(plain.data, ordinary.data, "requesting a twin must not change the ordinary path");
|
||||
let cached = cached.unwrap();
|
||||
assert_eq!((plain.width, plain.height), (24, 32));
|
||||
assert_eq!((cached.width, cached.height), (plain.width, plain.height));
|
||||
let worst = plain.data.iter().zip(&cached.data).flat_map(|(a, b)| a.iter().zip(b).map(|(a, b)| (a - b).abs())).fold(0.0f32, f32::max);
|
||||
assert!(worst < 0.002, "color finish differs beyond cache half-float rounding: {worst}");
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
/// Luminance (from linear RGB) minus its 9 × 9 local mean: tone differences between a render and a camera JPEG
|
||||
/// mostly cancel, edges remain.
|
||||
fn detail(w: usize, h: usize, lum: impl Fn(usize, usize) -> f32) -> Vec<f32> {
|
||||
@@ -676,8 +822,12 @@ mod tests {
|
||||
assert!(info.camera_tone.is_none());
|
||||
// a map that removes all saturation, and a tone curve
|
||||
let grey = HsvTable { hue_divisions: 4, sat_divisions: 2, val_divisions: 1, data: vec![[0.0, 0.0, 1.0]; 8], srgb_value: false };
|
||||
raw.color.profile =
|
||||
ProfileLook { hue_sat_map: [Some(grey), None], look_table: None, tone_curve: ToneCurve::from_tag(&[0.0, 0.0, 0.18, 0.3, 1.0, 1.0]) };
|
||||
raw.color.profile = ProfileLook {
|
||||
hue_sat_map: [Some(grey), None],
|
||||
look_table: None,
|
||||
tone_curve: ToneCurve::from_tag(&[0.0, 0.0, 0.18, 0.3, 1.0, 1.0]),
|
||||
gain_table_map: None,
|
||||
};
|
||||
let with = lightcraft_raw::write_dng(&raw, &Default::default()).unwrap();
|
||||
let (after, info) = load_bytes(&with, 64).unwrap();
|
||||
assert!(sat(&after) < 1e-3, "saturation {} → {}", sat(&before), sat(&after));
|
||||
@@ -695,6 +845,119 @@ mod tests {
|
||||
assert!((mean(&dim) / mean(&before) - 0.5).abs() < 0.02, "{} vs {}", mean(&dim), mean(&before));
|
||||
}
|
||||
|
||||
/// iPhone ProRAW-style semantic masks drive the Sky mask: Apple's sky matte (here the right
|
||||
/// half of the sensor, i.e. the bottom of the photo once turned 90° clockwise — where the sky
|
||||
/// heuristic would never look) is used; a matte with an unknown name is ignored.
|
||||
#[test]
|
||||
fn dng_sky_matte_drives_the_sky_mask() {
|
||||
use lightcraft_develop::{DevelopSettings, LocalAdjustments, Mask, MaskComponent, MaskOp, MaskShape};
|
||||
use lightcraft_pipeline::{RenderRequest, render};
|
||||
use lightcraft_tiff::tags::{self as t, photometric};
|
||||
use lightcraft_tiff::{ByteOrder, IfdBuilder, ImageData, TiffWriter, Value};
|
||||
let (w, h) = (16u32, 8u32);
|
||||
let image = |ifd: &mut IfdBuilder, w: u32, h: u32, cpp: u16, bits: u16, data: Vec<u8>| {
|
||||
ifd.set(t::IMAGE_WIDTH, Value::Long(vec![w]));
|
||||
ifd.set(t::IMAGE_LENGTH, Value::Long(vec![h]));
|
||||
ifd.set(t::BITS_PER_SAMPLE, Value::Short(vec![bits; cpp as usize]));
|
||||
ifd.set(t::SAMPLES_PER_PIXEL, Value::Short(vec![cpp]));
|
||||
ifd.set(t::COMPRESSION, Value::Short(vec![1]));
|
||||
ifd.set_image(ImageData::Strips { rows_per_strip: h, strips: vec![data] });
|
||||
};
|
||||
let mut raw = IfdBuilder::new();
|
||||
raw.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![0]));
|
||||
raw.set(t::PHOTOMETRIC, Value::Short(vec![photometric::LINEAR_RAW]));
|
||||
image(&mut raw, w, h, 3, 16, 12000u16.to_le_bytes().repeat((w * h * 3) as usize));
|
||||
let matte = |name: &str, right: bool| {
|
||||
let mut m = IfdBuilder::new();
|
||||
m.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![t::SUBFILE_SEMANTIC_MASK]));
|
||||
m.set(t::PHOTOMETRIC, Value::Short(vec![photometric::MASK]));
|
||||
m.set(t::SEMANTIC_NAME, Value::Ascii(name.into()));
|
||||
image(&mut m, w / 2, h / 2, 1, 8, (0..w / 2 * h / 2).map(|i| if (i % (w / 2) >= w / 4) == right { 255 } else { 0 }).collect());
|
||||
m
|
||||
};
|
||||
let mut ifd0 = IfdBuilder::new();
|
||||
ifd0.set(t::NEW_SUBFILE_TYPE, Value::Long(vec![1]));
|
||||
ifd0.set(t::DNG_VERSION, Value::Byte(vec![1, 6, 0, 0]));
|
||||
ifd0.set(t::MAKE, Value::Ascii("Apple".into()));
|
||||
ifd0.set(t::ORIENTATION, Value::Short(vec![6]));
|
||||
ifd0.set(t::COLOR_MATRIX_1, Value::SRational(vec![(1, 1), (0, 1), (0, 1), (0, 1), (1, 1), (0, 1), (0, 1), (0, 1), (1, 1)]));
|
||||
ifd0.set(t::CALIBRATION_ILLUMINANT_1, Value::Short(vec![21]));
|
||||
ifd0.set(t::PHOTOMETRIC, Value::Short(vec![photometric::RGB]));
|
||||
image(&mut ifd0, 4, 2, 3, 8, vec![128; 24]);
|
||||
ifd0.add_sub_ifd(raw);
|
||||
ifd0.add_sub_ifd(matte("urn:com:apple:photo:2020:aux:semanticskymatte", true));
|
||||
ifd0.add_sub_ifd(matte("urn:com:apple:photo:2020:aux:semanticsomethingelse", false));
|
||||
let bytes = TiffWriter::new(ByteOrder::Little, false).write(&[ifd0]).unwrap();
|
||||
|
||||
let (img, info) = load_bytes(&bytes, 64).unwrap();
|
||||
assert_eq!((img.width, img.height), (8, 16));
|
||||
assert!(info.mattes.is_some());
|
||||
let sky = Mask {
|
||||
components: vec![MaskComponent { name: None, op: MaskOp::Add, invert: false, shape: MaskShape::Sky }],
|
||||
adjust: LocalAdjustments { exposure: -3.0, ..Default::default() },
|
||||
..Default::default()
|
||||
};
|
||||
let s = DevelopSettings { masks: vec![sky], ..Default::default() };
|
||||
let req = RenderRequest::fit(8, 16);
|
||||
let green = |s: &DevelopSettings, info: &SourceInfo, y: usize| render(&img, info, s, &req).image.get(4, y)[1] as i32;
|
||||
let plain = DevelopSettings::default();
|
||||
assert_eq!(green(&s, &info, 2), green(&plain, &info, 2), "no sky at the top");
|
||||
assert!(green(&s, &info, 13) + 40 < green(&plain, &info, 13), "the sky at the bottom is darkened");
|
||||
// without the matte the heuristic looks at the top of the frame instead
|
||||
let heuristic = SourceInfo { mattes: None, ..info.clone() };
|
||||
assert!(green(&s, &heuristic, 2) < green(&plain, &heuristic, 2));
|
||||
assert_eq!(green(&s, &heuristic, 13), green(&plain, &heuristic, 13));
|
||||
}
|
||||
|
||||
/// A DNG's white balance is re-evaluated through its own colour model: the decoder hands the
|
||||
/// pipeline the file's colour tags (without the profile look, applied at load) and the white
|
||||
/// the pixels were developed for, and a custom white balance renders a neutral differently
|
||||
/// from the Bradford adaptation it replaces.
|
||||
#[test]
|
||||
fn dng_white_balance_uses_the_camera_colour_model() {
|
||||
use lightcraft_pipeline::{RenderRequest, render};
|
||||
let dng = crate::tests_xmp::synthetic_dng_with(None, Default::default());
|
||||
let raw = lightcraft_raw::decode(&dng).unwrap();
|
||||
let (img, info) = load_bytes(&dng, 64).unwrap();
|
||||
let cc = info.camera_color.as_ref().expect("a DNG with a colour matrix");
|
||||
assert_eq!(cc.developed_for, lightcraft_raw::color::as_shot_white_xy(&raw));
|
||||
assert!(cc.tags.profile.is_empty());
|
||||
assert_eq!(cc.tags.color_matrix, raw.color.color_matrix);
|
||||
let mut s = lightcraft_develop::DevelopSettings::default();
|
||||
(s.wb.mode, s.wb.temp, s.wb.tint) = (lightcraft_develop::WbMode::Custom, 3000.0, 20.0);
|
||||
let req = RenderRequest::fit(32, 32);
|
||||
let camera = render(&img, &info, &s, &req).image;
|
||||
let adapted = render(&img, &SourceInfo { camera_color: None, ..info.clone() }, &s, &req).image;
|
||||
assert_ne!(camera.data, adapted.data);
|
||||
}
|
||||
|
||||
/// Apple ProRAW: like Lightroom Classic, the default render ignores the file's
|
||||
/// `ProfileGainTableMap` (Apple's local tone mapping); a map that would double every pixel
|
||||
/// changes nothing.
|
||||
#[test]
|
||||
fn dng_gain_table_map_is_not_rendered() {
|
||||
use lightcraft_raw::gaintable::GainTableMap;
|
||||
let plain = crate::tests_xmp::synthetic_dng_with(None, Default::default());
|
||||
let mut raw = lightcraft_raw::decode(&plain).unwrap();
|
||||
let (before, _) = load_bytes(&lightcraft_raw::write_dng(&raw, &Default::default()).unwrap(), 64).unwrap();
|
||||
raw.color.profile.gain_table_map = Some(GainTableMap {
|
||||
points_v: 1,
|
||||
points_h: 1,
|
||||
points_n: 1,
|
||||
spacing_v: 1.0,
|
||||
spacing_h: 1.0,
|
||||
origin_v: 0.0,
|
||||
origin_h: 0.0,
|
||||
weights: [0.2, 0.2, 0.2, 0.2, 0.2],
|
||||
gamma: 1.0,
|
||||
gains: vec![2.0],
|
||||
});
|
||||
let with_map = lightcraft_raw::write_dng(&raw, &Default::default()).unwrap();
|
||||
assert!(lightcraft_raw::decode(&with_map).unwrap().color.profile.gain_table_map.is_some(), "the map is kept");
|
||||
let (after, _) = load_bytes(&with_map, 64).unwrap();
|
||||
assert_eq!(before.data, after.data);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_raw_falls_back_to_embedded_preview() {
|
||||
let b = cr3_with_preview(48, 32);
|
||||
|
||||
@@ -10,15 +10,26 @@
|
||||
//! run them off the UI thread.
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
// Model/cache types remain available without linking the optional inference crate.
|
||||
#[cfg(not(feature = "denoise"))]
|
||||
extern crate lightcraft_denoise_core as lightcraft_denoise;
|
||||
|
||||
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;
|
||||
pub mod denoise;
|
||||
pub mod devices;
|
||||
pub mod export;
|
||||
pub mod face_download;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod faces_index;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod faces_worker;
|
||||
pub mod files;
|
||||
pub mod fonts;
|
||||
pub mod guard;
|
||||
@@ -33,6 +44,7 @@ pub mod logging;
|
||||
pub mod media;
|
||||
pub mod memory;
|
||||
pub mod merge;
|
||||
mod model_download;
|
||||
pub mod originals;
|
||||
pub mod preset_import;
|
||||
pub mod preset_luminar;
|
||||
@@ -132,8 +144,10 @@ pub struct Interaction {
|
||||
|
||||
/// Source of [`Session::visible_shared`] generations (process-wide, so two sessions never share one).
|
||||
static VISIBLE_GEN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
|
||||
static LIBRARY_GEN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
|
||||
|
||||
pub struct Session {
|
||||
library_generation: u64,
|
||||
/// Auto Sync: edits to the active photo also change the other selected photos (the settings
|
||||
/// that changed, nothing else).
|
||||
pub auto_sync: bool,
|
||||
@@ -157,6 +171,17 @@ 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,
|
||||
/// AI denoise: the model in use, the photos that have their picture and the work in progress.
|
||||
pub(crate) denoise: denoise::State,
|
||||
/// 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>,
|
||||
/// Face model downloads started this session (the staged files wait in `<face_models_dir>/.downloads`).
|
||||
pub face_downloads: face_download::Downloads,
|
||||
/// The user's own list of models to download (`catalog.json` in the models folder), as last read.
|
||||
pub(crate) face_catalog: lightcraft_faces::catalog::Catalog,
|
||||
/// The loaded recognition model and the face embeddings made with it.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
pub(crate) faces: faces_index::FacesState,
|
||||
/// Copied develop settings (partial JSON) for Paste.
|
||||
pub clipboard: Option<Value>,
|
||||
/// The folder on disk the [`LibrarySource::Folder`] view browses.
|
||||
@@ -247,8 +272,13 @@ impl Default for Session {
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub fn library_generation(&self) -> u64 {
|
||||
self.library_generation
|
||||
}
|
||||
|
||||
pub fn new() -> Session {
|
||||
Session {
|
||||
library_generation: LIBRARY_GEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
|
||||
auto_sync: false,
|
||||
catalog: Catalog::new(),
|
||||
source: LibrarySource::All,
|
||||
@@ -263,6 +293,12 @@ impl Session {
|
||||
redo: Vec::new(),
|
||||
interaction: None,
|
||||
skip_auto_write: false,
|
||||
denoise: Default::default(),
|
||||
face_models_dir: None,
|
||||
face_downloads: Default::default(),
|
||||
face_catalog: Default::default(),
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
faces: Default::default(),
|
||||
clipboard: None,
|
||||
meta_clipboard: None,
|
||||
browse: None,
|
||||
@@ -403,6 +439,27 @@ impl Session {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The box to cut a face's picture from: the detector's, when the scan has looked at the face (every face is then shown
|
||||
/// equally close, however loosely or tightly its own region was drawn), else the region's own box.
|
||||
pub fn face_view(&self, id: PhotoId, rect: lightcraft_geom::Rect) -> lightcraft_geom::Rect {
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
if let Some(v) = self.faces.index.view(id.0, &rect) {
|
||||
return v;
|
||||
}
|
||||
let _ = id;
|
||||
rect
|
||||
}
|
||||
|
||||
/// Apply an op that is LightCraft's own bookkeeping rather than something the user did (faces found by the background
|
||||
/// scan): journaled like any op, but not an undo step, and it leaves the redo stack alone.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
pub(crate) fn apply_system(&mut self, op: Op) -> Result<()> {
|
||||
let fwd = op.clone();
|
||||
self.catalog.apply(op)?;
|
||||
self.pending_log.push(fwd);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Leaving photo `id` after editing it: keep its settings as an automatic version (when they
|
||||
/// differ from its latest version; at most [`AUTO_VERSIONS`] auto versions, oldest dropped).
|
||||
/// Saved with the library but not an undo step.
|
||||
@@ -781,10 +838,19 @@ pub fn json_delta(old: &Value, new: &Value) -> Option<Value> {
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
#[cfg(test)]
|
||||
mod tests_album_order;
|
||||
#[cfg(test)]
|
||||
mod tests_color;
|
||||
#[cfg(test)]
|
||||
mod tests_denoise;
|
||||
#[cfg(test)]
|
||||
mod tests_export;
|
||||
#[cfg(test)]
|
||||
mod tests_face_models;
|
||||
#[cfg(test)]
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod tests_face_recognize;
|
||||
#[cfg(test)]
|
||||
mod tests_folders;
|
||||
#[cfg(test)]
|
||||
mod tests_forget_local;
|
||||
|
||||
@@ -297,6 +297,12 @@ impl Session {
|
||||
{
|
||||
log::error!("library: {e}");
|
||||
}
|
||||
// Photo ids belong to one library; denoise work must not carry into the next.
|
||||
self.denoise.library_changed();
|
||||
self.media.denoise.clear();
|
||||
// what was learned about the old library's faces is saved, and none of it carries over (photo ids are per library)
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
self.faces.library_changed();
|
||||
let (mut journal, mut catalog, report) = Journal::open(catalog)?;
|
||||
// A loaded catalog counts revisions from 0, like every other one. Caches (sidebar counts,
|
||||
// keyword tree, the grid's list…) are keyed on the revision, so give each library loaded
|
||||
@@ -305,6 +311,7 @@ impl Session {
|
||||
static LOADS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
|
||||
catalog.revision = LOADS.fetch_add(1, std::sync::atomic::Ordering::Relaxed).wrapping_add(1) << 32;
|
||||
self.catalog = catalog;
|
||||
self.library_generation = crate::LIBRARY_GEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
self.undo.clear();
|
||||
self.redo.clear();
|
||||
self.interaction = None;
|
||||
|
||||
+248
-30
@@ -62,7 +62,7 @@ impl SettingsHashes {
|
||||
}
|
||||
|
||||
/// Bump when the pipeline's output changes, to invalidate cached thumbnails.
|
||||
pub const RENDER_CACHE_VERSION: u64 = 15;
|
||||
pub const RENDER_CACHE_VERSION: u64 = 16;
|
||||
|
||||
/// Thumbnails render at one of these long edges (so window/cell size changes reuse the cache).
|
||||
pub const THUMB_SIZES: [usize; 4] = [128, 256, 384, 512];
|
||||
@@ -125,6 +125,43 @@ pub type FileLoader = Arc<dyn Fn(&str, usize) -> Result<(Rgb32f, SourceInfo), St
|
||||
/// `max_edge` (set by the app). `None`: no usable preview.
|
||||
pub type PreviewLoader = Arc<dyn Fn(&str, usize) -> Option<Rgba8> + Send + Sync>;
|
||||
|
||||
/// Makes a photo's denoised picture when it is not cached yet (an export does not wait for the background queue):
|
||||
/// `Ok(true)` once it is there, `Ok(false)` when denoise does not apply to this photo (render it plain), `Err` when it
|
||||
/// should have worked and did not.
|
||||
pub type MakeProduct = Arc<dyn Fn() -> Result<bool, String> + Send + Sync>;
|
||||
|
||||
/// A raw photo's cached AI-denoised picture (see [`crate::denoise`]): where the file is and the key it must have been
|
||||
/// made with. A file that is missing, damaged or made for something else is a cache miss, never an error.
|
||||
#[derive(Clone)]
|
||||
pub struct DenoiseSpec {
|
||||
pub product: std::path::PathBuf,
|
||||
pub key: String,
|
||||
/// Run before the picture is read, by a job that must not render without it (an export).
|
||||
pub make: Option<MakeProduct>,
|
||||
}
|
||||
|
||||
/// Decodes a raw file once and develops it twice: as usual, and from the denoised picture in `spec` when that can be
|
||||
/// read (`None` when it cannot). Both are at most `max_edge` long and the same size.
|
||||
pub type PairLoader = Arc<dyn Fn(&str, usize, &DenoiseSpec) -> Result<(Rgb32f, Option<Rgb32f>, SourceInfo), String> + Send + Sync>;
|
||||
|
||||
/// A photo's denoised picture next to its plain one, and the last mix of the two that was asked for.
|
||||
pub struct Twin {
|
||||
pub image: Arc<Rgb32f>,
|
||||
/// The Amount (in thousandths) and the picture mixed for it, so dragging other sliders keeps one source picture.
|
||||
mix: std::sync::Mutex<Option<(u32, Arc<Rgb32f>)>>,
|
||||
}
|
||||
|
||||
impl Twin {
|
||||
pub fn new(image: Arc<Rgb32f>) -> Twin {
|
||||
Twin { image, mix: std::sync::Mutex::new(None) }
|
||||
}
|
||||
|
||||
/// Bytes reserved for the twin and its lazy mix, before either cache cost can grow.
|
||||
fn bytes(&self) -> usize {
|
||||
self.image.data.len().saturating_mul(24)
|
||||
}
|
||||
}
|
||||
|
||||
/// Pixels and the source interpretation learned while decoding; they must be cached together.
|
||||
#[derive(Clone)]
|
||||
pub struct DecodedSource {
|
||||
@@ -133,17 +170,52 @@ pub struct DecodedSource {
|
||||
/// A smart preview's stored camera tone curve: the one decoder fact its pixels need that the
|
||||
/// catalog's header facts lack (used when `info` is `None`).
|
||||
pub camera_tone: Option<lightcraft_pipeline::tone::CameraTone>,
|
||||
/// The same picture developed from the AI-denoised mosaic, when the photo has one cached.
|
||||
pub denoised: Option<Arc<Twin>>,
|
||||
}
|
||||
|
||||
impl DecodedSource {
|
||||
pub fn new(image: Arc<Rgb32f>, info: Option<SourceInfo>) -> Self {
|
||||
DecodedSource { image, info, camera_tone: None }
|
||||
DecodedSource { image, info, camera_tone: None, denoised: None }
|
||||
}
|
||||
|
||||
/// The picture to develop for a Denoise `amount` (0 to 1, see [`DevelopSettings::denoise_amount`]): the plain one,
|
||||
/// the denoised one, or a mix of the two. A mix is made once per amount and kept, so the stages after it are
|
||||
/// reused until the amount changes.
|
||||
pub fn image_for(&self, amount: f32) -> Arc<Rgb32f> {
|
||||
let Some(twin) = &self.denoised else { return self.image.clone() };
|
||||
if amount.is_nan() || amount <= 0.0 || twin.image.width != self.image.width || twin.image.height != self.image.height {
|
||||
return self.image.clone();
|
||||
}
|
||||
if amount >= 1.0 {
|
||||
return twin.image.clone();
|
||||
}
|
||||
let thousandths = (amount * 1000.0).round().clamp(1.0, 999.0) as u32;
|
||||
let mut memo = twin.mix.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some((t, img)) = memo.as_ref()
|
||||
&& *t == thousandths
|
||||
{
|
||||
return img.clone();
|
||||
}
|
||||
let mut mixed = (*self.image).clone();
|
||||
if !lightcraft_denoise::product::blend(&mut mixed, &twin.image, thousandths as f32 / 1000.0) {
|
||||
return self.image.clone();
|
||||
}
|
||||
let mixed = Arc::new(mixed);
|
||||
*memo = Some((thousandths, mixed.clone()));
|
||||
mixed
|
||||
}
|
||||
|
||||
/// Cache cost: the picture, its segmentation mattes, its denoised twin and room for the lazy mix.
|
||||
pub fn bytes(&self) -> usize {
|
||||
let mattes = self.info.as_ref().and_then(|i| i.mattes.as_ref()).map_or(0, |m| m.bytes());
|
||||
self.image.data.len() * 12 + 64 + std::mem::size_of::<SourceInfo>() + mattes + self.denoised.as_ref().map_or(0, |t| t.bytes())
|
||||
}
|
||||
|
||||
/// What to render these pixels against: the decoder's facts, else `header` (the catalog's)
|
||||
/// with any stored camera tone curve.
|
||||
pub fn info_or(&self, header: SourceInfo) -> SourceInfo {
|
||||
self.info.unwrap_or(SourceInfo { camera_tone: self.camera_tone.or(header.camera_tone), ..header })
|
||||
self.info.clone().unwrap_or_else(|| SourceInfo { camera_tone: self.camera_tone.or(header.camera_tone), ..header })
|
||||
}
|
||||
}
|
||||
|
||||
@@ -161,6 +233,8 @@ pub enum SourceRef {
|
||||
/// The photo's smart preview, used when the original can't be read (offline drive). The
|
||||
/// UI thread never checks whether the original is there; the render worker finds out.
|
||||
fallback: Option<std::path::PathBuf>,
|
||||
/// The photo has a cached denoised picture: load it with the plain one (see [`PairLoader`]).
|
||||
denoise: Option<(DenoiseSpec, PairLoader)>,
|
||||
},
|
||||
/// A smart preview standing in for a missing original.
|
||||
Smart {
|
||||
@@ -177,15 +251,45 @@ impl SourceRef {
|
||||
let image = match self {
|
||||
SourceRef::Loaded(a) => return Ok((**a).clone()),
|
||||
SourceRef::Demo { scene, max_edge } => crate::demo::scene_pixels(scene, *max_edge),
|
||||
SourceRef::File { path, max_edge, loader, fallback } => {
|
||||
let r = match loader {
|
||||
SourceRef::File { path, max_edge, loader, fallback, denoise } => {
|
||||
let plain = || match loader {
|
||||
Some(l) => l(path, *max_edge).map(|(image, info)| DecodedSource::new(Arc::new(image), Some(info))),
|
||||
None => Err(format!("no decoder available for {path}")),
|
||||
};
|
||||
// the denoised picture comes with the plain one from a single decode; when it cannot be read the
|
||||
// pair loader gives the plain picture alone (the cached file is never a reason to fail a render)
|
||||
let r = match denoise {
|
||||
Some((spec, pair)) => {
|
||||
// a job that needs the denoised picture has it made first; if that fails the job fails (no smart
|
||||
// preview stands in: an export must not quietly come out without the denoise it was asked for)
|
||||
let made = match &spec.make {
|
||||
Some(make) => make()?,
|
||||
None => true,
|
||||
};
|
||||
if made {
|
||||
pair(path, *max_edge, spec).and_then(|(image, twin, info)| {
|
||||
if spec.make.is_some() && twin.is_none() {
|
||||
return Err("the cached AI Denoise picture could not be read; clear the denoise cache and retry".into());
|
||||
}
|
||||
Ok(DecodedSource {
|
||||
image: Arc::new(image),
|
||||
info: Some(info),
|
||||
camera_tone: None,
|
||||
denoised: twin.map(|t| Arc::new(Twin::new(Arc::new(t)))),
|
||||
})
|
||||
})
|
||||
} else {
|
||||
plain()
|
||||
}
|
||||
}
|
||||
None => plain(),
|
||||
};
|
||||
// An offline original renders from its smart preview (no decoder facts: header ones).
|
||||
return match (r, fallback) {
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
(Err(e), Some(sp)) if crate::smart::is_valid(sp) => crate::smart::load(sp).map_err(|e2| format!("{e}; smart preview: {e2}")),
|
||||
(Err(e), Some(sp)) if denoise.as_ref().is_none_or(|(spec, _)| spec.make.is_none()) && crate::smart::is_valid(sp) => {
|
||||
crate::smart::load(sp).map_err(|e2| format!("{e}; smart preview: {e2}"))
|
||||
}
|
||||
(r, _) => r,
|
||||
};
|
||||
}
|
||||
@@ -224,6 +328,54 @@ pub struct MediaCache {
|
||||
scenes: Vec<lightcraft_scenes::Scene>,
|
||||
/// Rendered thumbnails (memory, plus disk once a library is attached).
|
||||
pub rendered: Arc<PreviewCache>,
|
||||
/// Which photos have a cached denoised picture, and how to read it.
|
||||
pub denoise: DenoiseIndex,
|
||||
}
|
||||
|
||||
/// The photos that have a cached denoised picture for the denoise model in use. Only memory: nothing here touches the
|
||||
/// file system, so it can be asked on the UI thread.
|
||||
#[derive(Default)]
|
||||
pub struct DenoiseIndex {
|
||||
ready: std::collections::HashMap<PhotoId, DenoiseSpec>,
|
||||
/// Reads the pictures (`None`: this build cannot, and nothing is ever ready).
|
||||
pub loader: Option<PairLoader>,
|
||||
}
|
||||
|
||||
impl DenoiseIndex {
|
||||
/// What to load for `id` besides its plain picture, and how.
|
||||
pub fn get(&self, id: PhotoId) -> Option<(&DenoiseSpec, &PairLoader)> {
|
||||
self.ready.get(&id).zip(self.loader.as_ref())
|
||||
}
|
||||
|
||||
pub fn spec(&self, id: PhotoId) -> Option<&DenoiseSpec> {
|
||||
self.ready.get(&id)
|
||||
}
|
||||
|
||||
pub fn set(&mut self, id: PhotoId, spec: DenoiseSpec) {
|
||||
self.ready.insert(id, spec);
|
||||
}
|
||||
|
||||
pub fn remove(&mut self, id: PhotoId) -> bool {
|
||||
self.ready.remove(&id).is_some()
|
||||
}
|
||||
|
||||
pub fn clear(&mut self) {
|
||||
self.ready.clear();
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.ready.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.ready.is_empty()
|
||||
}
|
||||
|
||||
pub fn ids(&self) -> Vec<PhotoId> {
|
||||
let mut v: Vec<PhotoId> = self.ready.keys().copied().collect();
|
||||
v.sort();
|
||||
v
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MediaCache {
|
||||
@@ -244,6 +396,7 @@ impl Default for MediaCache {
|
||||
availability: Default::default(),
|
||||
scenes: Vec::new(),
|
||||
rendered: Arc::new(PreviewCache::memory(rendered_budget(budget))),
|
||||
denoise: DenoiseIndex::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -253,10 +406,6 @@ fn rendered_budget(share: usize) -> usize {
|
||||
RENDERED_MEM_BYTES.min(share / 4)
|
||||
}
|
||||
|
||||
fn source_bytes(img: &Rgb32f) -> usize {
|
||||
img.data.len() * 12 + 64 + std::mem::size_of::<SourceInfo>()
|
||||
}
|
||||
|
||||
impl MediaCache {
|
||||
/// Keep rendered thumbnails on disk in `dir` as well. The cache this replaces is retired
|
||||
/// (its files stay): render jobs still holding it can't write through it any more, even
|
||||
@@ -319,8 +468,8 @@ impl MediaCache {
|
||||
let Some((tick, which)) = candidates.into_iter().flatten().min() else { break };
|
||||
let freed = match which {
|
||||
0 => self.thumbs.pop_oldest(),
|
||||
1 => self.previews.iter().position(|e| e.2 == tick).map(|i| source_bytes(&self.previews.remove(i).1.image)),
|
||||
2 => self.full.take().map(|e| source_bytes(&e.1.image)),
|
||||
1 => self.previews.iter().position(|e| e.2 == tick).map(|i| self.previews.remove(i).1.bytes()),
|
||||
2 => self.full.take().map(|e| e.1.bytes()),
|
||||
_ => self.rendered.evict_oldest(),
|
||||
};
|
||||
match freed {
|
||||
@@ -372,7 +521,7 @@ impl MediaCache {
|
||||
let tick = lightcraft_preview::next_tick();
|
||||
match level {
|
||||
SourceLevel::Thumb => {
|
||||
let cost = source_bytes(&img.image);
|
||||
let cost = img.bytes();
|
||||
self.thumbs.insert(id, img, cost);
|
||||
}
|
||||
SourceLevel::Preview => {
|
||||
@@ -415,13 +564,42 @@ impl MediaCache {
|
||||
/// Decoded sources held: (thumbnail level, preview level, full size).
|
||||
pub fn usage(&self) -> (crate::memory::Usage, crate::memory::Usage, crate::memory::Usage) {
|
||||
use crate::memory::Usage;
|
||||
let previews = Usage::new(self.previews.len(), self.previews.iter().map(|e| source_bytes(&e.1.image)).sum());
|
||||
let full = self.full.as_ref().map(|e| Usage::new(1, source_bytes(&e.1.image))).unwrap_or_default();
|
||||
let previews = Usage::new(self.previews.len(), self.previews.iter().map(|e| e.1.bytes()).sum());
|
||||
let full = self.full.as_ref().map(|e| Usage::new(1, e.1.bytes())).unwrap_or_default();
|
||||
(Usage::new(self.thumbs.len(), self.thumbs.cost()), previews, full)
|
||||
}
|
||||
|
||||
/// What a render of photo `id` with settings `s` adds to its identity for the denoised picture it will use: 0 unless
|
||||
/// it uses one (the Denoise amount is above 0 and the photo has a cached denoised picture).
|
||||
pub fn denoise_salt(&self, id: PhotoId, s: &DevelopSettings) -> u64 {
|
||||
match self.denoise.spec(id) {
|
||||
Some(spec) if s.denoise_amount() > 0.0 => Hasher128::new().str(&spec.key).finish().0 as u64 | 1,
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// How to load photo `p` at `level`: from the cache when its source is there, else from its file (with its denoised
|
||||
/// picture beside it when it has one), its smart preview or the demo scene.
|
||||
pub fn source_ref(&mut self, p: &Photo, level: SourceLevel) -> SourceRef {
|
||||
if let Some(a) = self.get_source(p.id, level) {
|
||||
self.source_ref_with(p, level, true)
|
||||
}
|
||||
|
||||
/// [`Self::source_ref`] that never takes a decoded source from memory: for a job that needs something the cached
|
||||
/// one lacks (an export, whose denoised picture is made when it runs).
|
||||
pub fn source_ref_uncached(&mut self, p: &Photo, level: SourceLevel) -> SourceRef {
|
||||
self.source_ref_with(p, level, false)
|
||||
}
|
||||
|
||||
fn source_ref_with(&mut self, p: &Photo, level: SourceLevel, cached: bool) -> SourceRef {
|
||||
let mut r = self.source_ref_plain(p, level, cached);
|
||||
if let (SourceRef::File { denoise, .. }, Some((spec, loader))) = (&mut r, self.denoise.get(p.id)) {
|
||||
*denoise = Some((spec.clone(), loader.clone()));
|
||||
}
|
||||
r
|
||||
}
|
||||
|
||||
fn source_ref_plain(&mut self, p: &Photo, level: SourceLevel, cached: bool) -> SourceRef {
|
||||
if cached && let Some(a) = self.get_source(p.id, level) {
|
||||
return SourceRef::Loaded(Box::new(a));
|
||||
}
|
||||
// Procedural scenes have a nominal size: "full" is that size, not unbounded.
|
||||
@@ -459,10 +637,12 @@ impl MediaCache {
|
||||
}
|
||||
match self.scenes.iter().find(|s| s.id == *scene) {
|
||||
Some(s) => SourceRef::Demo { scene: Box::new(s.clone()), max_edge },
|
||||
None => SourceRef::File { path: format!("demo:{scene}"), max_edge, loader: None, fallback: None },
|
||||
None => SourceRef::File { path: format!("demo:{scene}"), max_edge, loader: None, fallback: None, denoise: None },
|
||||
}
|
||||
}
|
||||
Source::File { path } => SourceRef::File { path: path.clone(), max_edge, loader: self.file_loader.clone(), fallback: None },
|
||||
Source::File { path } => {
|
||||
SourceRef::File { path: path.clone(), max_edge, loader: self.file_loader.clone(), fallback: None, denoise: None }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -573,7 +753,8 @@ impl RenderJob {
|
||||
let was_loaded = matches!(self.source, SourceRef::Loaded(_));
|
||||
match self.source.load_source() {
|
||||
Ok(source) => {
|
||||
let src = &source.image;
|
||||
// the plain picture, or (Denoise amount above 0 and a denoised picture cached) the mix the amount asks for
|
||||
let src = &source.image_for(self.settings.denoise_amount());
|
||||
let info = source.info_or(self.info);
|
||||
// Thumbnails (many small jobs side by side) stay on the CPU; views and exports use
|
||||
// the GPU when there is one.
|
||||
@@ -757,22 +938,23 @@ impl crate::Session {
|
||||
let content = source_key.0 as u64;
|
||||
// Temporary "before" settings must not replace the photo's cached develop hash.
|
||||
let settings_hash = if before { settings.hash64() } else { self.media.settings_hashes.get(id, &settings) };
|
||||
let salt = self.media.denoise_salt(id, &settings);
|
||||
let key = settings_hash
|
||||
^ ((max_w as u64) << 40)
|
||||
^ ((max_h as u64) << 20)
|
||||
^ (apply_crop as u64)
|
||||
^ (level as u64) << 60
|
||||
^ id.0.wrapping_mul(0x9e37_79b9_7f4a_7c15)
|
||||
^ content.rotate_left(17);
|
||||
^ content.rotate_left(17)
|
||||
^ salt;
|
||||
let cache = thumb_bucket.map(|b| {
|
||||
let k = Hasher128::new()
|
||||
.str(&content_key)
|
||||
.u64(settings_hash)
|
||||
.u64(b as u64)
|
||||
.u64(RENDER_CACHE_VERSION)
|
||||
.u64(crate::camera_profiles::cache_key())
|
||||
.finish();
|
||||
(self.media.rendered.clone(), k)
|
||||
let mut k = Hasher128::new();
|
||||
k.str(&content_key).u64(settings_hash).u64(b as u64).u64(RENDER_CACHE_VERSION).u64(crate::camera_profiles::cache_key());
|
||||
// (only a render that uses a denoised picture has it in its key, so every other thumbnail stays valid)
|
||||
if salt != 0 {
|
||||
k.u64(salt);
|
||||
}
|
||||
(self.media.rendered.clone(), k.finish())
|
||||
});
|
||||
Some(RenderJob {
|
||||
request_id: 0,
|
||||
@@ -979,6 +1161,17 @@ impl crate::Session {
|
||||
}
|
||||
}
|
||||
|
||||
/// The camera's embedded preview of a raw file, for the face scan. Faces are read from the unedited picture, so the
|
||||
/// photo's edits do not matter, and the camera's own JPEG (oriented, display-ready) costs a fraction of decoding the
|
||||
/// raw. `None` for anything else, and where the host installed no loader.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
pub(crate) fn scan_preview_of(&self, p: &Photo) -> Option<(String, PreviewLoader)> {
|
||||
match (&p.source, &self.media.preview_loader) {
|
||||
(Source::File { path }, Some(l)) if p.kind == MediaKind::Raw => Some((path.clone(), l.clone())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Something to show in the loupe right away for `id` (see [`QuickJob`]): its cached view
|
||||
/// render, else the embedded preview of an unedited raw, else a cached or fresh thumbnail.
|
||||
pub fn quick_view_job(&mut self, id: PhotoId, max_edge: usize, apply_crop: bool) -> Option<QuickJob> {
|
||||
@@ -1052,6 +1245,15 @@ impl crate::Session {
|
||||
depth: lightcraft_pipeline::OutputDepth,
|
||||
) -> Result<RenderJob, String> {
|
||||
let mut job = self.render_job(id, max_w, max_h, false, true).ok_or("no such photo")?;
|
||||
// An export is made from the denoised picture its Denoise amount asks for, made when the job runs if the
|
||||
// background queue has not got to it. (A source kept in memory may lack the picture, so the job reads the file.)
|
||||
if let Some(wanted) = self.denoise_for_export(id, &job.settings)? {
|
||||
let p = self.catalog.photo(id).ok_or("no such photo")?.clone();
|
||||
job.source = self.media.source_ref_uncached(&p, job.level);
|
||||
if let SourceRef::File { denoise, .. } = &mut job.source {
|
||||
*denoise = Some(wanted);
|
||||
}
|
||||
}
|
||||
job.request.space = space;
|
||||
job.request.depth = depth;
|
||||
job.key ^= (space as u64 + 1).wrapping_mul(0xa076_1d64_78bd_642f) ^ (depth as u64 + 1).wrapping_mul(0xe703_7ed1_a0b4_28db);
|
||||
@@ -1199,6 +1401,20 @@ mod thumbnail_hash_tests {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn denoise_mix_is_reserved_before_the_source_enters_the_cache() {
|
||||
let mut plain = Rgb32f::new(8, 8);
|
||||
plain.data.fill([1.0; 3]);
|
||||
let mut source = DecodedSource::new(Arc::new(plain), None);
|
||||
source.denoised = Some(Arc::new(Twin::new(Arc::new(Rgb32f::new(8, 8)))));
|
||||
let reserved = source.bytes();
|
||||
assert!(reserved >= 3 * 8 * 8 * 12);
|
||||
assert_eq!(source.image_for(0.5).data[0], [0.5; 3]);
|
||||
assert_eq!(source.bytes(), reserved, "creating a mix cannot evade the LRU's recorded cost");
|
||||
let _ = source.image_for(0.25);
|
||||
assert_eq!(source.bytes(), reserved, "changing Amount replaces the reserved mix");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_info_survives_render_jobs_cache_and_eviction() {
|
||||
let tone = lightcraft_pipeline::tone::CameraTone::new(std::array::from_fn(|i| {
|
||||
@@ -1210,20 +1426,22 @@ mod tests {
|
||||
let mut s = crate::Session::with_demo();
|
||||
let id = s.active().unwrap();
|
||||
let mut job = s.render_job(id, 64, 64, false, true).unwrap();
|
||||
let loaded_info = info.clone();
|
||||
job.source = SourceRef::File {
|
||||
path: "synthetic.arw".into(),
|
||||
max_edge: 64,
|
||||
loader: Some(Arc::new(move |_, _| {
|
||||
let mut image = Rgb32f::new(64, 64);
|
||||
image.data.fill([0.1; 3]);
|
||||
Ok((image, info))
|
||||
Ok((image, loaded_info.clone()))
|
||||
})),
|
||||
fallback: None,
|
||||
denoise: None,
|
||||
};
|
||||
job.info = SourceInfo::default(); // Header facts cannot override decoder facts.
|
||||
job.settings = Arc::new(DevelopSettings::default());
|
||||
let r = job.clone().run();
|
||||
assert_eq!(r.loaded.as_ref().unwrap().info, Some(info));
|
||||
assert_eq!(r.loaded.as_ref().unwrap().info.as_ref(), Some(&info));
|
||||
let expected = lightcraft_pipeline::render(&r.loaded.as_ref().unwrap().image, &info, &job.settings, &job.request);
|
||||
assert_eq!(r.rendered.as_ref().unwrap().image.data, expected.image.data);
|
||||
s.accept(&r);
|
||||
|
||||
@@ -0,0 +1,359 @@
|
||||
//! Downloading a denoise model explicitly requested after showing its terms.
|
||||
//!
|
||||
//! The transfer is [`lightcraft_fetch`] (pure-Rust HTTPS with resume, size and SHA-256 checks), on a background thread,
|
||||
//! so the UI never waits for the network. Only a model listed in [`lightcraft_denoise::known::Download`] can be fetched:
|
||||
//! the address and hash come from our pinned catalog or the user's model catalog, and no token is sent with it. What arrives
|
||||
//! matches the model's recorded size and SHA-256 or it is thrown away; it then waits in a staging folder until the model
|
||||
//! is installed (`denoise.models.install`, or by itself once the user has accepted the model's terms in the download).
|
||||
//!
|
||||
//! Failure is a message, never a panic: no network, a stalled or truncated transfer, a full disk, a file that is not the
|
||||
//! one expected, and a panic inside the thread itself all end as [`State::Failed`].
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex, PoisonError};
|
||||
use std::time::Duration;
|
||||
|
||||
use lightcraft_denoise::known::Download;
|
||||
|
||||
/// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it
|
||||
/// can never be a model's own folder (ids start with a letter or digit).
|
||||
pub const STAGING: &str = ".downloads";
|
||||
|
||||
/// Where a download stands.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum State {
|
||||
/// Fetching (`bytes` of `total` so far); the file is checked as the last step.
|
||||
Running {
|
||||
bytes: u64,
|
||||
total: u64,
|
||||
},
|
||||
/// Fetched and verified (its SHA-256 is `sha256`), waiting in the staging folder at `path` to be installed.
|
||||
Done {
|
||||
path: PathBuf,
|
||||
sha256: String,
|
||||
},
|
||||
/// Installed and in use: nothing more to do but tell the user.
|
||||
Installing,
|
||||
Installed,
|
||||
Failed(String),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
state: Mutex<State>,
|
||||
cancel: AtomicBool,
|
||||
}
|
||||
|
||||
impl Shared {
|
||||
fn new(total: u64) -> Self {
|
||||
Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) }
|
||||
}
|
||||
|
||||
fn set(&self, s: State) {
|
||||
*self.state.lock().unwrap_or_else(PoisonError::into_inner) = s;
|
||||
}
|
||||
|
||||
fn get(&self) -> State {
|
||||
self.state.lock().unwrap_or_else(PoisonError::into_inner).clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// The downloads of one session, by model id.
|
||||
#[derive(Default)]
|
||||
pub struct Downloads {
|
||||
jobs: HashMap<String, Arc<Shared>>,
|
||||
}
|
||||
|
||||
impl Downloads {
|
||||
/// Start fetching `spec` into `<models dir>/.downloads/`. Different models can download side by side; the same
|
||||
/// one is started once.
|
||||
pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> {
|
||||
if self.jobs.get(&spec.id).is_some_and(|j| matches!(j.get(), State::Running { .. })) {
|
||||
return Err("that model is already downloading".into());
|
||||
}
|
||||
let staging = models_dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?;
|
||||
let shared = Arc::new(Shared::new(spec.size_bytes));
|
||||
let worker = shared.clone();
|
||||
let id = spec.id.clone();
|
||||
std::thread::Builder::new()
|
||||
.name("denoise-model-download".into())
|
||||
.spawn(move || {
|
||||
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, &spec.url, &staging, &worker)));
|
||||
if outcome.is_err() {
|
||||
worker.set(State::Failed("the download stopped unexpectedly".into()));
|
||||
}
|
||||
})
|
||||
.map_err(|e| format!("could not start the download: {e}"))?;
|
||||
self.jobs.insert(id, shared);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Every download this session has started and not discarded, by model id.
|
||||
pub fn snapshot(&self) -> Vec<(String, State)> {
|
||||
let mut all: Vec<_> = self.jobs.iter().map(|(id, j)| (id.clone(), j.get())).filter(|(_, s)| *s != State::Cancelled).collect();
|
||||
all.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
all
|
||||
}
|
||||
|
||||
pub fn running(&self) -> bool {
|
||||
self.jobs.values().any(|j| matches!(j.get(), State::Running { .. }))
|
||||
}
|
||||
|
||||
/// Stop a running download, or discard a finished one (its staged file is deleted). `true` if there was one.
|
||||
pub fn discard(&mut self, id: &str) -> bool {
|
||||
let Some(job) = self.jobs.get(id) else { return false };
|
||||
job.cancel.store(true, Ordering::Relaxed);
|
||||
match job.get() {
|
||||
// the thread notices, deletes the partial file and marks itself cancelled
|
||||
State::Running { .. } | State::Installing => {}
|
||||
State::Done { path, .. } => {
|
||||
let _ = std::fs::remove_file(path);
|
||||
self.jobs.remove(id);
|
||||
}
|
||||
State::Installed | State::Failed(_) | State::Cancelled => {
|
||||
self.jobs.remove(id);
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Downloads that have arrived and been checked, ready to install: (model id, staged file, its SHA-256).
|
||||
pub fn finished(&self) -> Vec<(String, PathBuf, String)> {
|
||||
let mut out: Vec<_> = self
|
||||
.jobs
|
||||
.iter()
|
||||
.filter_map(|(id, j)| match j.get() {
|
||||
State::Done { path, sha256 } => Some((id.clone(), path, sha256)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
out.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
out
|
||||
}
|
||||
|
||||
/// Record how installing a finished download went (`Installed` or `Failed`).
|
||||
pub fn set_outcome(&mut self, id: &str, state: State) {
|
||||
match self.jobs.get(id) {
|
||||
Some(j) => j.set(state),
|
||||
None => {
|
||||
let j = Shared::new(0);
|
||||
j.set(state);
|
||||
self.jobs.insert(id.to_string(), Arc::new(j));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Downloads {
|
||||
/// Quitting while a download runs stops it instead of leaving it behind.
|
||||
fn drop(&mut self) {
|
||||
for job in self.jobs.values() {
|
||||
job.cancel.store(true, Ordering::Relaxed);
|
||||
}
|
||||
let started = web_time::Instant::now();
|
||||
while self.running() && started.elapsed() < Duration::from_secs(1) {
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What went wrong, in words the user can act on.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
fn explain(e: &lightcraft_fetch::DownloadError) -> String {
|
||||
use lightcraft_fetch::DownloadError;
|
||||
match e {
|
||||
DownloadError::AllMirrorsFailed { errors, .. } => {
|
||||
let detail: String = errors.join("; ").chars().take(300).collect();
|
||||
format!("could not download the model ({detail}): check the internet connection, or get the file from its page and add it yourself")
|
||||
}
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch `spec` from its exact URL into `staging`, check it and leave it there.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
fn fetch(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) -> Result<(), lightcraft_fetch::DownloadError> {
|
||||
let file = lightcraft_fetch::FileSpec { name: &spec.file_name, size: Some(spec.size_bytes), sha256: Some(&spec.sha256), max: spec.size_bytes };
|
||||
lightcraft_fetch::download_url(&file, mirror, staging, &lightcraft_fetch::Options::default(), &shared.cancel, &mut |p| {
|
||||
shared.set(State::Running { bytes: p.done.min(spec.size_bytes), total: spec.size_bytes });
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
fn fetch(_: &Download, _: &str, _: &Path, _: &Shared) -> Result<(), String> {
|
||||
Err("the web build downloads no models".into())
|
||||
}
|
||||
|
||||
/// The whole job: reuse a verified earlier copy, else fetch, check and stage. Always ends in a final [`State`].
|
||||
fn run(spec: &Download, mirror: &str, staging: &Path, shared: &Shared) {
|
||||
let target = staging.join(&spec.file_name);
|
||||
let part = staging.join(format!("{}.part", spec.file_name));
|
||||
match fetch(spec, mirror, staging, shared) {
|
||||
Ok(()) => shared.set(State::Done { path: target, sha256: spec.sha256.to_string() }),
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
Err(lightcraft_fetch::DownloadError::Cancelled) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
shared.set(State::Cancelled);
|
||||
}
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
Err(e) => {
|
||||
shared.set(State::Failed(explain(&e)));
|
||||
}
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
Err(e) => {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
shared.set(State::Failed(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, not(target_arch = "wasm32")))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use lightcraft_denoise::hash::sha256_hex;
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::net::TcpListener;
|
||||
|
||||
fn scratch(name: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
fn spec_for(bytes: &[u8]) -> Download {
|
||||
Download {
|
||||
id: "test-model".into(),
|
||||
url: "https://example.invalid/m/test-model.onnx".into(),
|
||||
file_name: "test-model.onnx".into(),
|
||||
size_bytes: bytes.len() as u64,
|
||||
sha256: sha256_hex(bytes),
|
||||
}
|
||||
}
|
||||
|
||||
/// A local server that answers every request with `status` and `body` (the whole file, no ranges).
|
||||
fn serve(status: &'static str, body: Vec<u8>) -> String {
|
||||
let l = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let base = format!("http://{}/m", l.local_addr().unwrap());
|
||||
std::thread::spawn(move || {
|
||||
for mut s in l.incoming().flatten() {
|
||||
let mut r = BufReader::new(s.try_clone().unwrap());
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = s.write_all(format!("HTTP/1.1 {status}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes());
|
||||
let _ = s.write_all(&body);
|
||||
}
|
||||
});
|
||||
base
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fetches_checks_and_stages_a_file() {
|
||||
let dir = scratch("ok");
|
||||
let bytes: Vec<u8> = (0..200_000u32).map(|i| (i % 251) as u8).collect();
|
||||
let (staging, spec) = (dir.join(STAGING), spec_for(&bytes));
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let mirror = serve("200 OK", bytes.clone());
|
||||
let shared = Shared::new(spec.size_bytes);
|
||||
run(&spec, &mirror, &staging, &shared);
|
||||
let State::Done { path, sha256 } = shared.get() else { panic!("{:?}", shared.get()) };
|
||||
assert_eq!((std::fs::read(&path).unwrap(), sha256), (bytes, spec.sha256.to_string()));
|
||||
assert!(!staging.join("test-model.onnx.part").exists());
|
||||
// asked again with the copy still staged: no fetch needed (nothing listens there)
|
||||
let again = Shared::new(spec.size_bytes);
|
||||
run(&spec, "http://127.0.0.1:1/m", &staging, &again);
|
||||
assert!(matches!(again.get(), State::Done { .. }), "{:?}", again.get());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_file_that_does_not_match_is_thrown_away() {
|
||||
let dir = scratch("bad");
|
||||
let staging = dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
// the right size and the wrong content
|
||||
let spec = spec_for(&[7u8; 5000]);
|
||||
let shared = Shared::new(spec.size_bytes);
|
||||
run(&spec, &serve("200 OK", vec![8u8; 5000]), &staging, &shared);
|
||||
assert!(matches!(shared.get(), State::Failed(_)), "{:?}", shared.get());
|
||||
// nothing is left behind: not the file, not a partial one
|
||||
assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unreachable_or_refusing_server_is_a_message() {
|
||||
let dir = scratch("down");
|
||||
let staging = dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let spec = spec_for(b"abc");
|
||||
for mirror in ["http://127.0.0.1:1/m".to_string(), serve("404 Not Found", b"not found".to_vec())] {
|
||||
let shared = Shared::new(3);
|
||||
run(&spec, &mirror, &staging, &shared);
|
||||
let State::Failed(why) = shared.get() else { panic!("{:?}", shared.get()) };
|
||||
assert!(why.contains("internet connection"), "{why}");
|
||||
}
|
||||
assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_cancelled_download_leaves_nothing() {
|
||||
let dir = scratch("cancel");
|
||||
let staging = dir.join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let spec = spec_for(b"abc");
|
||||
let shared = Shared::new(3);
|
||||
shared.cancel.store(true, Ordering::Relaxed);
|
||||
run(&spec, &serve("200 OK", b"abc".to_vec()), &staging, &shared);
|
||||
assert_eq!(shared.get(), State::Cancelled);
|
||||
assert_eq!(std::fs::read_dir(&staging).unwrap().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dropping_the_session_stops_a_running_download() {
|
||||
let shared = Arc::new(Shared::new(10));
|
||||
// a stand-in for the download thread: notices the cancel flag the way the downloader does
|
||||
let worker = shared.clone();
|
||||
let handle = std::thread::spawn(move || {
|
||||
while !worker.cancel.load(Ordering::Relaxed) {
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
worker.set(State::Cancelled);
|
||||
});
|
||||
let mut d = Downloads::default();
|
||||
d.jobs.insert("m".into(), shared.clone());
|
||||
assert!(d.running());
|
||||
let started = web_time::Instant::now();
|
||||
drop(d);
|
||||
assert!(started.elapsed() < Duration::from_millis(900), "quitting must not wait for a transfer to finish");
|
||||
assert_eq!(shared.get(), State::Cancelled);
|
||||
handle.join().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discarding_a_finished_download_deletes_its_file() {
|
||||
let dir = scratch("discard");
|
||||
let file = dir.join("staged.onnx");
|
||||
std::fs::write(&file, b"x").unwrap();
|
||||
let mut d = Downloads::default();
|
||||
let shared = Arc::new(Shared::new(1));
|
||||
shared.set(State::Done { path: file.clone(), sha256: String::new() });
|
||||
d.jobs.insert("m".into(), shared);
|
||||
assert_eq!(d.snapshot().len(), 1);
|
||||
assert!(!d.running());
|
||||
assert!(d.discard("m"));
|
||||
assert!(!file.exists());
|
||||
assert!(d.snapshot().is_empty());
|
||||
assert!(!d.discard("m"), "nothing left to discard");
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
}
|
||||
@@ -157,7 +157,14 @@ impl Segmenter {
|
||||
/// Whether this build can compute AI masks at all.
|
||||
pub const AVAILABLE: bool = cfg!(feature = "sam");
|
||||
|
||||
/// Whether the model's files are in place.
|
||||
pub fn remote_endpoint(&self) -> Option<String> {
|
||||
#[cfg(feature = "sam")]
|
||||
return lightcraft_segment::remote::configured();
|
||||
#[cfg(not(feature = "sam"))]
|
||||
None
|
||||
}
|
||||
|
||||
/// Whether the local model files are in place.
|
||||
pub fn installed(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
if let Some(dir) = &self.dir {
|
||||
@@ -230,6 +237,10 @@ impl Segmenter {
|
||||
if !Self::AVAILABLE {
|
||||
return Err("AI masks are not available in this build".into());
|
||||
}
|
||||
#[cfg(feature = "sam")]
|
||||
if lightcraft_segment::remote::configured().is_some() {
|
||||
return Ok(PathBuf::new());
|
||||
}
|
||||
let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?;
|
||||
if self.installed() {
|
||||
return Ok(dir);
|
||||
|
||||
@@ -17,7 +17,10 @@ use std::sync::mpsc::{self, RecvTimeoutError};
|
||||
use std::time::Duration;
|
||||
|
||||
use lightcraft_develop::SegMask;
|
||||
use lightcraft_segment::{Encoded, Sam3};
|
||||
use lightcraft_segment::{
|
||||
Encoded, Sam3,
|
||||
remote::{self, Operation, Request, text_logits},
|
||||
};
|
||||
|
||||
use super::{Input, Prompt, Tag};
|
||||
use crate::media::RenderJob;
|
||||
@@ -175,6 +178,7 @@ impl Worker {
|
||||
struct State {
|
||||
model: Option<(PathBuf, Sam3)>,
|
||||
cache: Option<(u64, Encoded)>,
|
||||
remote: Option<remote::Client>,
|
||||
}
|
||||
|
||||
/// Whether `later` makes `job` pointless (it asks for the same selection again).
|
||||
@@ -195,7 +199,7 @@ fn run(rx: &mpsc::Receiver<Job>, shared: &Shared, idle: Duration) {
|
||||
match rx.recv_timeout(idle) {
|
||||
Ok(j) => queue.push_back(j),
|
||||
Err(RecvTimeoutError::Timeout) => {
|
||||
if state.model.is_some() {
|
||||
if state.model.is_some() || state.remote.is_some() {
|
||||
state = State::default();
|
||||
shared.loaded.store(false, Ordering::SeqCst);
|
||||
log::info!("SAM 3 unloaded after {} minutes without use", idle.as_secs() / 60);
|
||||
@@ -231,7 +235,7 @@ fn run(rx: &mpsc::Receiver<Job>, shared: &Shared, idle: Duration) {
|
||||
if result.is_err() && !is_detail {
|
||||
state.cache = None;
|
||||
}
|
||||
shared.loaded.store(state.model.is_some(), Ordering::SeqCst);
|
||||
shared.loaded.store(state.model.is_some() || state.remote.is_some(), Ordering::SeqCst);
|
||||
drop(done);
|
||||
let _ = job.reply.send(Outcome { tag: job.tag, result, superseded: false });
|
||||
}
|
||||
@@ -284,29 +288,68 @@ fn encoded<'a>(state: &'a mut State, shared: &Shared, job: &mut Job) -> Result<(
|
||||
log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed());
|
||||
state.cache = Some((job.key, enc));
|
||||
}
|
||||
let State { model, cache } = state;
|
||||
let State { model, cache, .. } = state;
|
||||
match (model.as_mut(), cache.as_mut()) {
|
||||
(Some((_, m)), Some((_, e))) => Ok((m, e)),
|
||||
_ => Err("the model is not ready".into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every phrase of `text` ("car, road"), merged (the per-pixel maximum); `None` when nothing
|
||||
/// matches.
|
||||
fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result<Option<Vec<f32>>, String> {
|
||||
let mut merged: Option<Vec<f32>> = None;
|
||||
for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) {
|
||||
if let Some(p) = model.segment_text(enc, phrase, 0.5).map_err(|e| e.to_string())? {
|
||||
merged = Some(match merged {
|
||||
None => p.logits,
|
||||
Some(m) => m.iter().zip(&p.logits).map(|(a, b)| a.max(*b)).collect(),
|
||||
});
|
||||
}
|
||||
fn remote_job(state: &mut State, shared: &Shared, job: &mut Job, endpoint: &str) -> Result<Option<SegMask>, String> {
|
||||
let _busy = Flag::raise(&shared.encoding);
|
||||
if state.remote.is_none() {
|
||||
state.remote = Some(remote::Client::connect(endpoint)?);
|
||||
}
|
||||
Ok(merged)
|
||||
let client = state.remote.as_mut().ok_or("remote SAM unavailable")?;
|
||||
let mut region = None;
|
||||
let operation = match std::mem::replace(&mut job.kind, Kind::Prepare) {
|
||||
Kind::Prepare => Operation::Prepare,
|
||||
Kind::Text(text) => Operation::Text(text),
|
||||
Kind::Clicks(clicks) => Operation::Clicks(clicks.iter().map(|c| [c.x, c.y, if c.positive { 1.0 } else { 0.0 }]).collect()),
|
||||
Kind::Detail { prompt, region: r } => {
|
||||
region = Some(r);
|
||||
let (rw, rh) = (r[2] - r[0], r[3] - r[1]);
|
||||
if !(rw > 0.0 && rh > 0.0) {
|
||||
return Ok(None);
|
||||
}
|
||||
match prompt {
|
||||
Prompt::Text(text) => Operation::Text(text),
|
||||
Prompt::Clicks(clicks) => {
|
||||
let inside: Vec<_> = clicks
|
||||
.iter()
|
||||
.map(|c| [((c.at.x - r[0]) / rw) as f32, ((c.at.y - r[1]) / rh) as f32, if c.include { 1.0 } else { 0.0 }])
|
||||
.filter(|p| (0.0..=1.0).contains(&p[0]) && (0.0..=1.0).contains(&p[1]))
|
||||
.collect();
|
||||
if !inside.iter().any(|p| p[2] == 1.0) {
|
||||
return Ok(None);
|
||||
}
|
||||
Operation::Clicks(inside)
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
Kind::Panic => return Err("test panic request".into()),
|
||||
};
|
||||
let input = if region.is_some() || client.key != Some(job.key) { render(job)? } else { Input { rgb: Vec::new(), w: 0, h: 0 } };
|
||||
let request = Request { width: input.w, height: input.h, temporary: region.is_some(), operation };
|
||||
let logits = client.request(&request, &input.rgb)?;
|
||||
if region.is_none() {
|
||||
client.key = Some(job.key);
|
||||
}
|
||||
Ok(logits.map(|l| match region {
|
||||
Some(r) => SegMask::from_logits_in(lightcraft_segment::MASK_SIDE, &l, r),
|
||||
None => SegMask::from_logits(lightcraft_segment::MASK_SIDE, &l),
|
||||
}))
|
||||
}
|
||||
|
||||
fn run_job(state: &mut State, shared: &Shared, job: &mut Job) -> Result<Option<SegMask>, String> {
|
||||
if let Some(endpoint) = remote::configured() {
|
||||
let result = remote_job(state, shared, job, &endpoint);
|
||||
if result.is_err() {
|
||||
state.remote = None;
|
||||
}
|
||||
return result.map_err(|e| format!("Remote SAM 3 failed: {e}. Check the worker and SSH tunnel, then retry."));
|
||||
}
|
||||
let side = lightcraft_segment::MASK_SIDE;
|
||||
match std::mem::replace(&mut job.kind, Kind::Prepare) {
|
||||
Kind::Prepare => {
|
||||
|
||||
@@ -190,7 +190,7 @@ pub fn is_valid(path: &Path) -> bool {
|
||||
/// Load the proxy at `path`.
|
||||
pub fn load(path: &Path) -> Result<crate::media::DecodedSource, String> {
|
||||
let b = std::fs::read(path).map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
decode(&b).map(|(image, camera_tone)| crate::media::DecodedSource { image: Arc::new(image), info: None, camera_tone })
|
||||
decode(&b).map(|(image, camera_tone)| crate::media::DecodedSource { image: Arc::new(image), info: None, camera_tone, denoised: None })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -118,6 +118,29 @@ fn albums_crud() {
|
||||
s.execute("album.delete", &json!({"id": f})).unwrap();
|
||||
}
|
||||
|
||||
/// Given an album and a folder, when something is created inside each, then only the folder
|
||||
/// accepts it: the sidebar only lists children of folders, so a child of a plain album would
|
||||
/// be invisible.
|
||||
#[test]
|
||||
fn albums_and_smart_albums_are_created_only_inside_folders() {
|
||||
let mut s = demo();
|
||||
let folder = s.execute("album.create", &json!({"name": "Trips", "folder": true})).unwrap()["id"].as_u64().unwrap();
|
||||
let album = s.execute("album.create", &json!({"name": "Best", "parent": folder})).unwrap()["id"].as_u64().unwrap();
|
||||
for (cmd, params) in [
|
||||
("album.create", json!({"name": "Nested", "parent": album})),
|
||||
("album.create", json!({"name": "Nested", "folder": true, "parent": album})),
|
||||
("album.createSmart", json!({"name": "Nested", "parent": album})),
|
||||
("album.create", json!({"name": "Nested", "parent": 987_654})),
|
||||
] {
|
||||
assert!(s.execute(cmd, ¶ms).is_err(), "{cmd} {params} must be refused");
|
||||
}
|
||||
let n = s.catalog.albums().count();
|
||||
let sub = s.execute("album.create", &json!({"name": "Sub", "folder": true, "parent": folder})).unwrap()["id"].as_u64().unwrap();
|
||||
let smart = s.execute("album.createSmart", &json!({"name": "Rated", "parent": sub})).unwrap()["id"].as_u64().unwrap();
|
||||
assert_eq!(s.catalog.albums().count(), n + 2);
|
||||
assert_eq!(s.catalog.album(lightcraft_catalog::AlbumId(smart)).unwrap().parent, Some(lightcraft_catalog::AlbumId(sub)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filters_and_delete_restore() {
|
||||
let mut s = demo();
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
//! Ordering albums by hand: `album.reorder` puts an album (or folder) before a sibling, or last of
|
||||
//! its kind, in one undo step, also when it changes folder; `album.sort` goes back to by name.
|
||||
|
||||
use lightcraft_catalog::AlbumId;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::Session;
|
||||
|
||||
fn make(s: &mut Session, name: &str, parent: Option<u64>, folder: bool) -> u64 {
|
||||
let mut p = json!({"name": name, "folder": folder});
|
||||
if let Some(parent) = parent {
|
||||
p["parent"] = json!(parent);
|
||||
}
|
||||
s.execute("album.create", &p).unwrap()["id"].as_u64().unwrap()
|
||||
}
|
||||
|
||||
fn names(s: &Session, parent: Option<u64>) -> Vec<String> {
|
||||
s.catalog.album_children(parent.map(AlbumId)).iter().map(|a| a.name.clone()).collect()
|
||||
}
|
||||
|
||||
fn reorder(s: &mut Session, id: u64, params: Value) -> Result<Value, String> {
|
||||
let mut p = params;
|
||||
p["id"] = json!(id);
|
||||
s.execute("album.reorder", &p).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
/// Given three albums listed by name, when one is placed before another, then the list follows;
|
||||
/// "no sibling" puts it last (among albums).
|
||||
#[test]
|
||||
fn an_album_is_placed_before_a_sibling_or_last() {
|
||||
let mut s = Session::new();
|
||||
let (a, b, c) = (make(&mut s, "A", None, false), make(&mut s, "B", None, false), make(&mut s, "C", None, false));
|
||||
reorder(&mut s, c, json!({"before": a})).unwrap();
|
||||
assert_eq!(names(&s, None), ["C", "A", "B"]);
|
||||
assert!(s.catalog.album_children_are_ordered(None));
|
||||
reorder(&mut s, c, json!({"before": null})).unwrap();
|
||||
assert_eq!(names(&s, None), ["A", "B", "C"]);
|
||||
reorder(&mut s, a, json!({})).unwrap();
|
||||
assert_eq!(names(&s, None), ["B", "C", "A"], "no `before`: last");
|
||||
assert_eq!(reorder(&mut s, b, json!({"before": b})).unwrap()["changed"], 0, "before itself is where it is");
|
||||
assert_eq!(names(&s, None), ["B", "C", "A"]);
|
||||
}
|
||||
|
||||
/// Folders stay ahead of albums: a folder is placed among folders, an album among albums.
|
||||
#[test]
|
||||
fn folders_and_albums_keep_to_their_own_group() {
|
||||
let mut s = Session::new();
|
||||
let (f1, f2) = (make(&mut s, "F1", None, true), make(&mut s, "F2", None, true));
|
||||
let (a1, a2) = (make(&mut s, "A1", None, false), make(&mut s, "A2", None, false));
|
||||
reorder(&mut s, f2, json!({"before": f1})).unwrap();
|
||||
reorder(&mut s, a2, json!({"before": a1})).unwrap();
|
||||
assert_eq!(names(&s, None), ["F2", "F1", "A2", "A1"]);
|
||||
reorder(&mut s, f1, json!({"before": a1})).expect_err("a folder is not placed among albums");
|
||||
reorder(&mut s, a1, json!({"before": f1})).expect_err("an album is not placed among folders");
|
||||
reorder(&mut s, f2, json!({"before": null})).unwrap();
|
||||
assert_eq!(names(&s, None), ["F1", "F2", "A2", "A1"], "last of the folders");
|
||||
}
|
||||
|
||||
/// Placing an album in another folder moves it there, and one undo undoes both.
|
||||
#[test]
|
||||
fn placing_across_folders_moves_it_and_undoes_in_one_step() {
|
||||
let mut s = Session::new();
|
||||
let folder = make(&mut s, "Trips", None, true);
|
||||
let (x, y) = (make(&mut s, "X", Some(folder), false), make(&mut s, "Y", Some(folder), false));
|
||||
let loose = make(&mut s, "Loose", None, false);
|
||||
reorder(&mut s, y, json!({"before": x})).unwrap();
|
||||
assert_eq!(names(&s, Some(folder)), ["Y", "X"]);
|
||||
reorder(&mut s, loose, json!({"parent": folder, "before": x})).unwrap();
|
||||
assert_eq!(names(&s, Some(folder)), ["Y", "Loose", "X"]);
|
||||
assert_eq!(s.catalog.album(AlbumId(loose)).unwrap().parent, Some(AlbumId(folder)));
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert_eq!(names(&s, Some(folder)), ["Y", "X"]);
|
||||
assert_eq!(s.catalog.album(AlbumId(loose)).unwrap().parent, None);
|
||||
assert_eq!(s.catalog.album(AlbumId(loose)).unwrap().order, None);
|
||||
}
|
||||
|
||||
/// Undo and redo of placing across folders, of sorting, and of moving a folder that has children.
|
||||
#[test]
|
||||
fn placing_and_sorting_undo_and_redo() {
|
||||
let mut s = Session::new();
|
||||
let top = make(&mut s, "Top", None, true);
|
||||
let child = make(&mut s, "Child", Some(top), false);
|
||||
let (a, b) = (make(&mut s, "A", None, false), make(&mut s, "B", None, false));
|
||||
reorder(&mut s, b, json!({"before": a})).unwrap();
|
||||
// a folder with children, placed inside another folder
|
||||
let other = make(&mut s, "Other", None, true);
|
||||
reorder(&mut s, top, json!({"parent": other})).unwrap();
|
||||
assert_eq!(names(&s, Some(other)), ["Top"]);
|
||||
assert_eq!(names(&s, Some(top)), ["Child"], "its children came along");
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert_eq!(s.catalog.album(AlbumId(top)).unwrap().parent, None);
|
||||
s.execute("edit.redo", &json!({})).unwrap();
|
||||
assert_eq!(s.catalog.album(AlbumId(top)).unwrap().parent, Some(AlbumId(other)));
|
||||
// sort, undo, redo
|
||||
s.execute("album.sort", &json!({})).unwrap();
|
||||
assert_eq!(names(&s, None), ["Other", "A", "B"]);
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert_eq!(names(&s, None), ["Other", "B", "A"], "the hand order is back");
|
||||
s.execute("edit.redo", &json!({})).unwrap();
|
||||
assert_eq!(names(&s, None), ["Other", "A", "B"]);
|
||||
assert_eq!(s.catalog.album(AlbumId(child)).unwrap().parent, Some(AlbumId(top)));
|
||||
}
|
||||
|
||||
/// Bad requests change nothing: a `before` that is not in the target folder, a target that is not
|
||||
/// a folder or lies inside the album, an unknown album.
|
||||
#[test]
|
||||
fn bad_placements_are_refused_and_change_nothing() {
|
||||
let mut s = Session::new();
|
||||
let f = make(&mut s, "F", None, true);
|
||||
let inner = make(&mut s, "Inner", Some(f), true);
|
||||
let (a, b) = (make(&mut s, "A", None, false), make(&mut s, "B", Some(f), false));
|
||||
let before = s.catalog.to_snapshot();
|
||||
reorder(&mut s, a, json!({"before": b})).expect_err("B is in F, not at the top level");
|
||||
reorder(&mut s, a, json!({"parent": b})).expect_err("a plain album is no folder");
|
||||
reorder(&mut s, f, json!({"parent": inner})).expect_err("a folder into itself");
|
||||
reorder(&mut s, 987_654, json!({})).expect_err("no such album");
|
||||
reorder(&mut s, a, json!({"before": 987_654})).expect_err("no such sibling");
|
||||
assert_eq!(s.catalog.to_snapshot(), before, "nothing changed");
|
||||
}
|
||||
|
||||
/// `album.sort` drops the hand order of one folder (and only that one): by name again.
|
||||
#[test]
|
||||
fn sorting_goes_back_to_by_name_for_one_folder() {
|
||||
let mut s = Session::new();
|
||||
let f = make(&mut s, "F", None, true);
|
||||
let (x, y) = (make(&mut s, "X", Some(f), false), make(&mut s, "Y", Some(f), false));
|
||||
let (a, b) = (make(&mut s, "A", None, false), make(&mut s, "B", None, false));
|
||||
reorder(&mut s, y, json!({"before": x})).unwrap();
|
||||
reorder(&mut s, b, json!({"before": a})).unwrap();
|
||||
assert_eq!(names(&s, Some(f)), ["Y", "X"]);
|
||||
assert_eq!(names(&s, None), ["F", "B", "A"]);
|
||||
let r = s.execute("album.sort", &json!({"parent": f})).unwrap();
|
||||
assert_eq!(r["changed"], 2);
|
||||
assert_eq!(names(&s, Some(f)), ["X", "Y"]);
|
||||
assert_eq!(names(&s, None), ["F", "B", "A"], "the top level keeps its order");
|
||||
assert!(!s.catalog.album_children_are_ordered(Some(AlbumId(f))));
|
||||
// nothing to do: no undo step is made
|
||||
let r = s.execute("album.sort", &json!({"parent": f})).unwrap();
|
||||
assert_eq!(r["changed"], 0);
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert_eq!(names(&s, Some(f)), ["Y", "X"], "the one undo step was the sort");
|
||||
let r = s.execute("album.sort", &json!({})).unwrap();
|
||||
assert_eq!(r["changed"], 3, "every child of the top level had a place: F, B and A");
|
||||
assert_eq!(names(&s, None), ["F", "A", "B"]);
|
||||
}
|
||||
|
||||
/// Moving an album with `album.move` out of an ordered folder leaves the old neighbours' order.
|
||||
#[test]
|
||||
fn moving_out_of_an_ordered_folder_keeps_the_others_in_place() {
|
||||
let mut s = Session::new();
|
||||
let f = make(&mut s, "F", None, true);
|
||||
let (x, y, z) = (make(&mut s, "X", Some(f), false), make(&mut s, "Y", Some(f), false), make(&mut s, "Z", Some(f), false));
|
||||
reorder(&mut s, z, json!({"before": x})).unwrap();
|
||||
assert_eq!(names(&s, Some(f)), ["Z", "X", "Y"]);
|
||||
s.execute("album.move", &json!({"id": x, "parent": null})).unwrap();
|
||||
assert_eq!(names(&s, Some(f)), ["Z", "Y"]);
|
||||
assert_eq!(names(&s, None), ["F", "X"]);
|
||||
let _ = y;
|
||||
}
|
||||
@@ -0,0 +1,759 @@
|
||||
//! AI denoise through the engine: installing a model, making a photo's denoised picture in the background and for an
|
||||
//! export, the Amount slider mixing it in, and the cache. A tiny stand-in model (it halves every sample) takes the
|
||||
//! place of a real one, which is enough to see what the machinery does with its answer; how well real models denoise is
|
||||
//! measured separately (`crates/denoise/tests/real_model.rs`).
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use lightcraft_denoise::manifest::{DenoiserManifest, Domain, Gain};
|
||||
use lightcraft_denoise::run::TileRunner;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::Session;
|
||||
use crate::denoise::{Loader, Model, PhotoState, RunOn};
|
||||
use crate::export::{ExportOptions, export_photo};
|
||||
use crate::tests_xmp::{synthetic_dng_of, temp_dir};
|
||||
|
||||
/// A model that makes every output pixel half of its cell's colour: a picture that is plainly not the plain one.
|
||||
struct Dim {
|
||||
runs: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl TileRunner for Dim {
|
||||
fn run(&self, input: &[f32]) -> Result<Vec<f32>, lightcraft_denoise::Error> {
|
||||
self.runs.fetch_add(1, Ordering::Relaxed);
|
||||
let t = ((input.len() / 4) as f64).sqrt() as usize;
|
||||
let side = 2 * t;
|
||||
let mut out = vec![0f32; 3 * side * side];
|
||||
for y in 0..side {
|
||||
for x in 0..side {
|
||||
let i = (y / 2) * t + x / 2;
|
||||
let (r, g, b) = (input[i], (input[t * t + i] + input[2 * t * t + i]) / 2.0, input[3 * t * t + i]);
|
||||
out[y * side + x] = r * 0.5;
|
||||
out[side * side + y * side + x] = g * 0.5;
|
||||
out[2 * side * side + y * side + x] = b * 0.5;
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
impl Model for Dim {
|
||||
fn runner(&self, _: RunOn, threads: usize) -> (&dyn TileRunner, usize) {
|
||||
(self, threads)
|
||||
}
|
||||
|
||||
fn self_test(&self, _: RunOn) -> Result<Value, String> {
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
}
|
||||
|
||||
fn dim_loader(runs: &Arc<AtomicUsize>) -> Loader {
|
||||
let runs = runs.clone();
|
||||
Arc::new(move |_, _| Ok(Arc::new(Dim { runs: runs.clone() }) as Arc<dyn Model>))
|
||||
}
|
||||
|
||||
fn manifest(id: &str) -> DenoiserManifest {
|
||||
DenoiserManifest {
|
||||
id: id.into(),
|
||||
name: "Stand-in".into(),
|
||||
version: "1".into(),
|
||||
licence: Default::default(),
|
||||
source: None,
|
||||
sha256: None,
|
||||
size_bytes: None,
|
||||
provenance: "made up for a test".into(),
|
||||
domain: Domain::BayerToRgb,
|
||||
tile: 64,
|
||||
overlap: 16,
|
||||
gain: Gain::None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The model's `.onnx` (not a real one: the stand-in loader never reads it) with its manifest beside it.
|
||||
fn model_files(dir: &Path, id: &str) -> PathBuf {
|
||||
let src = dir.join(format!("src-{id}"));
|
||||
std::fs::create_dir_all(&src).unwrap();
|
||||
std::fs::write(src.join("denoise-model.json"), serde_json::to_vec(&manifest(id)).unwrap()).unwrap();
|
||||
std::fs::write(src.join("model.onnx"), format!("not an onnx file, only a stand-in for {id}")).unwrap();
|
||||
src.join("model.onnx")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_installation_keeps_the_frame_pump_responsive_and_reuses_the_loaded_model() {
|
||||
use std::sync::{Mutex, mpsc};
|
||||
use std::time::Duration;
|
||||
let mut x = setup("background-install", false);
|
||||
let path = model_files(&x.dir, "stand-in");
|
||||
let (entered_tx, entered_rx) = mpsc::channel();
|
||||
let (release_tx, release_rx) = mpsc::channel();
|
||||
let release_rx = Mutex::new(release_rx);
|
||||
let runs = x.runs.clone();
|
||||
let loads = Arc::new(AtomicUsize::new(0));
|
||||
let worker_loads = loads.clone();
|
||||
x.s.denoise.loader = Arc::new(move |_, _| {
|
||||
worker_loads.fetch_add(1, Ordering::Relaxed);
|
||||
entered_tx.send(()).map_err(|e| e.to_string())?;
|
||||
release_rx.lock().unwrap().recv_timeout(Duration::from_secs(3)).map_err(|e| e.to_string())?;
|
||||
Ok(Arc::new(Dim { runs: runs.clone() }))
|
||||
});
|
||||
let info = x.s.execute("denoise.models.inspect", &json!({"path": path})).unwrap();
|
||||
assert_eq!(info["model"]["id"], "stand-in");
|
||||
assert_eq!(loads.load(Ordering::Relaxed), 0, "showing terms never executes the model");
|
||||
let r = x.s.execute("denoise.models.install", &json!({"path": path, "acknowledged": true, "background": true})).unwrap();
|
||||
assert_eq!(r["started"], "local-install");
|
||||
entered_rx.recv_timeout(Duration::from_secs(3)).unwrap();
|
||||
// The loader is blocked: the frame pump must still return, reporting installation in progress.
|
||||
let state = x.s.execute("denoise.models.downloads", &json!({})).unwrap();
|
||||
assert_eq!(state["downloads"][0]["state"], "installing");
|
||||
assert_eq!(x.s.execute("denoise.pump", &json!({"pace": "pause"})).unwrap()["active"], false);
|
||||
assert!(x.s.execute("denoise.models.install", &json!({"path": path, "acknowledged": true, "background": true})).is_err());
|
||||
release_tx.send(()).unwrap();
|
||||
for _ in 0..600 {
|
||||
let v = x.s.execute("denoise.models.downloads", &json!({})).unwrap();
|
||||
if v["downloads"][0]["state"] == "installed" {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
assert_eq!(x.s.execute("denoise.models.downloads", &json!({})).unwrap()["downloads"][0]["state"], "installed");
|
||||
let a = x.s.denoise.active.as_ref().unwrap();
|
||||
assert!(a.model.lock().unwrap().is_some(), "the worker's tested model is reused");
|
||||
assert_eq!(loads.load(Ordering::Relaxed), 1);
|
||||
drop(x.s);
|
||||
let _ = std::fs::remove_dir_all(x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_model_replacement_preserves_the_installed_model() {
|
||||
let mut x = setup("replacement", true);
|
||||
let installed = x.dir.join("denoise-models/stand-in/model.onnx");
|
||||
let original = std::fs::read(&installed).unwrap();
|
||||
let source = model_files(&x.dir, "stand-in");
|
||||
std::fs::write(&source, "a different broken model").unwrap();
|
||||
assert!(x.s.execute("denoise.models.install", &json!({"path": source, "acknowledged": true})).is_err());
|
||||
assert_eq!(std::fs::read(installed).unwrap(), original);
|
||||
assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["model"], "stand-in");
|
||||
drop(x.s);
|
||||
let _ = std::fs::remove_dir_all(x.dir);
|
||||
}
|
||||
|
||||
/// A Bayer DNG that is a smooth ramp, big enough for a 64-cell tile and a second one.
|
||||
fn ramp_dng(w: usize, h: usize, cfa: &str) -> Vec<u8> {
|
||||
let data: Vec<u16> = (0..w * h).map(|i| 600 + ((i % w) * 30 + (i / w) * 35) as u16).collect();
|
||||
synthetic_dng_of(w, h, cfa, data, None, Default::default())
|
||||
}
|
||||
|
||||
struct Setup {
|
||||
dir: PathBuf,
|
||||
s: Session,
|
||||
runs: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
/// A session with a library on disk holding one raw photo, and the stand-in model installed and chosen.
|
||||
fn setup(tag: &str, with_model: bool) -> Setup {
|
||||
let dir = temp_dir(tag);
|
||||
let originals = dir.join("originals");
|
||||
std::fs::create_dir_all(&originals).unwrap();
|
||||
std::fs::write(originals.join("ramp.dng"), ramp_dng(192, 160, "RGGB")).unwrap();
|
||||
let runs = Arc::new(AtomicUsize::new(0));
|
||||
let mut s = Session::new().with_fs();
|
||||
s.open_library(dir.join("library"), false).unwrap();
|
||||
s.execute("library.import", &json!({"paths": [originals.join("ramp.dng").to_string_lossy()]})).unwrap();
|
||||
let id = s.catalog.photos().next().unwrap().id;
|
||||
s.execute("library.select", &json!({"ids": [id.0]})).unwrap();
|
||||
s.set_denoise_models_dir(Some(dir.join("denoise-models")));
|
||||
s.denoise.loader = dim_loader(&runs);
|
||||
if with_model {
|
||||
let onnx = model_files(&dir, "stand-in");
|
||||
s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
}
|
||||
Setup { dir, s, runs }
|
||||
}
|
||||
|
||||
fn photo(s: &Session) -> lightcraft_catalog::PhotoId {
|
||||
s.catalog.photos().next().unwrap().id
|
||||
}
|
||||
|
||||
/// Pump until `n` photos are ready (the job runs on its own thread).
|
||||
fn pump_until_ready(s: &mut Session, n: u64) -> Value {
|
||||
for _ in 0..600 {
|
||||
let r = s.execute("denoise.pump", &json!({"pace": "full"})).unwrap();
|
||||
if r["ready"].as_u64() >= Some(n) {
|
||||
return r;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
}
|
||||
panic!("not ready: {}", s.execute("denoise.status", &json!({})).unwrap());
|
||||
}
|
||||
|
||||
fn mean(img: &lightcraft_raster::Rgba8) -> f64 {
|
||||
let sum: u64 = img.data.iter().map(|p| u64::from(p[0]) + u64::from(p[1]) + u64::from(p[2])).sum();
|
||||
sum as f64 / (img.data.len() * 3) as f64
|
||||
}
|
||||
|
||||
fn rendered_mean(s: &mut Session, id: lightcraft_catalog::PhotoId) -> f64 {
|
||||
mean(&s.render_now(id, 160, 160).unwrap().image)
|
||||
}
|
||||
|
||||
fn set_amount(s: &mut Session, amount: f64) {
|
||||
s.execute("develop.set", &json!({"control": "enhance.denoise", "value": amount})).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_model_is_installed_chosen_and_removed_with_the_terms_accepted() {
|
||||
let mut x = setup("model", false);
|
||||
let onnx = model_files(&x.dir, "stand-in");
|
||||
// the terms come first
|
||||
let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy()})).unwrap_err().to_string();
|
||||
assert!(e.contains("licence has not been accepted"), "{e}");
|
||||
let r = x.s.execute("denoise.models.list", &json!({})).unwrap();
|
||||
assert_eq!(r["model"], Value::Null);
|
||||
assert_eq!(r["runtime"], cfg!(feature = "denoise"));
|
||||
// a file with no manifest is not guessed at
|
||||
let stray = x.dir.join("stray.onnx");
|
||||
std::fs::write(&stray, b"who knows").unwrap();
|
||||
let e = x.s.execute("denoise.models.install", &json!({"path": stray.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("does not know this file"), "{e}");
|
||||
// installed, it is the model in use
|
||||
let r = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
assert_eq!((r["installed"]["id"].clone(), r["model"].clone()), (json!("stand-in"), json!("stand-in")), "{r}");
|
||||
assert!(r["installed"]["accepted"]["sha256"].as_str().is_some_and(|h| h.len() == 64), "the file's hash is recorded: {r}");
|
||||
assert!(x.s.denoise.active.is_some());
|
||||
assert_eq!(x.s.execute("denoise.status", &json!({})).unwrap()["model"]["id"], "stand-in");
|
||||
// it can be checked again, switched off and on, and removed
|
||||
assert_eq!(x.s.execute("denoise.models.test", &json!({"id": "stand-in"})).unwrap()["ok"], true);
|
||||
assert_eq!(x.s.execute("denoise.models.select", &json!({"id": null})).unwrap()["model"], Value::Null);
|
||||
assert!(x.s.denoise.active.is_none());
|
||||
assert!(x.s.execute("denoise.models.select", &json!({"id": "nothing"})).is_err());
|
||||
x.s.execute("denoise.models.select", &json!({"id": "stand-in"})).unwrap();
|
||||
assert!(x.s.denoise.active.is_some());
|
||||
let r = x.s.execute("denoise.models.remove", &json!({"id": "stand-in"})).unwrap();
|
||||
assert_eq!(r["model"], Value::Null);
|
||||
assert!(x.s.denoise.active.is_none());
|
||||
assert!(!x.dir.join("denoise-models").join("stand-in").exists());
|
||||
// ids that are not folder names, and folders that are not models, are never touched
|
||||
for bad in ["../library", "", "a/b", "settings.json"] {
|
||||
assert!(x.s.execute("denoise.models.remove", &json!({"id": bad})).is_err(), "{bad}");
|
||||
}
|
||||
assert!(x.dir.join("library").is_dir());
|
||||
// a model that fails its self-test is not left installed
|
||||
x.s.denoise.loader = Arc::new(|_, _| Err("it does not load".into()));
|
||||
let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("it does not load"), "{e}");
|
||||
assert!(!x.dir.join("denoise-models").join("stand-in").exists());
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_manifest_that_does_not_validate_or_name_the_file_is_refused() {
|
||||
let mut x = setup("manifest", false);
|
||||
let onnx = model_files(&x.dir, "stand-in");
|
||||
let json_path = onnx.with_file_name("denoise-model.json");
|
||||
// a tile the model cannot run on
|
||||
let mut m = manifest("stand-in");
|
||||
m.tile = 50;
|
||||
std::fs::write(&json_path, serde_json::to_vec(&m).unwrap()).unwrap();
|
||||
let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("tile"), "{e}");
|
||||
// a hash that is not the file's
|
||||
let mut m = manifest("stand-in");
|
||||
m.sha256 = Some("0".repeat(64));
|
||||
std::fs::write(&json_path, serde_json::to_vec(&m).unwrap()).unwrap();
|
||||
let e = x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("SHA-256"), "{e}");
|
||||
assert!(!x.dir.join("denoise-models").join("stand-in").join("model.onnx").exists(), "nothing half-installed is left");
|
||||
// an id that would leave the folder
|
||||
let mut m = manifest("../escape");
|
||||
m.id = "../escape".into();
|
||||
std::fs::write(&json_path, serde_json::to_vec(&m).unwrap()).unwrap();
|
||||
assert!(x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).is_err());
|
||||
assert!(!x.dir.join("escape").exists());
|
||||
// garbage instead of a manifest
|
||||
std::fs::write(&json_path, b"{ not json").unwrap();
|
||||
assert!(x.s.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).is_err());
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_photo_is_denoised_in_the_background_and_the_amount_mixes_it_in() {
|
||||
let mut x = setup("background", true);
|
||||
let id = photo(&x.s);
|
||||
// nothing is made until a photo has an amount
|
||||
let r = x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap();
|
||||
assert_eq!((r["queued"].clone(), r["running"].clone(), r["ready"].clone()), (json!(0), Value::Null, json!(0)), "{r}");
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle);
|
||||
let plain = rendered_mean(&mut x.s, id);
|
||||
let key_before = x.s.render_job(id, 160, 160, false, true).unwrap().key;
|
||||
|
||||
set_amount(&mut x.s, 100.0);
|
||||
// the photo being edited is picked up by itself, and made on another thread
|
||||
let r = pump_until_ready(&mut x.s, 1);
|
||||
assert_eq!(r["queued"], 0, "{r}");
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready);
|
||||
assert!(x.runs.load(Ordering::Relaxed) >= 4, "a 192 × 160 mosaic is four 64-cell tiles: {}", x.runs.load(Ordering::Relaxed));
|
||||
let products: Vec<_> = std::fs::read_dir(x.dir.join("library").join("denoise")).unwrap().flatten().collect();
|
||||
assert_eq!(products.len(), 1, "one cached picture, no DNG or anything else in the library");
|
||||
assert!(products[0].path().extension().is_some_and(|e| e == "lcdn"));
|
||||
let originals: Vec<_> = std::fs::read_dir(x.dir.join("originals")).unwrap().flatten().collect();
|
||||
assert_eq!(originals.len(), 1, "the original is untouched and alone");
|
||||
|
||||
// the render now comes from the denoised picture: darker at 100, in between at 50, the plain one at 0
|
||||
let full = rendered_mean(&mut x.s, id);
|
||||
assert!(full < plain - 3.0, "denoised {full} vs plain {plain}");
|
||||
set_amount(&mut x.s, 50.0);
|
||||
let half = rendered_mean(&mut x.s, id);
|
||||
assert!(half < plain - 1.0 && half > full + 1.0, "{full} < {half} < {plain}");
|
||||
set_amount(&mut x.s, 0.0);
|
||||
assert!((rendered_mean(&mut x.s, id) - plain).abs() < 0.01, "amount 0 is the plain picture again");
|
||||
// …and a render that uses it has another identity than one that does not
|
||||
assert_ne!(key_before, {
|
||||
set_amount(&mut x.s, 100.0);
|
||||
x.s.render_job(id, 160, 160, false, true).unwrap().key
|
||||
});
|
||||
// switching the Detail section off switches the denoise off with it
|
||||
set_amount(&mut x.s, 100.0);
|
||||
x.s.execute("develop.sectionEnabled", &json!({"section": "detail", "enabled": false})).unwrap();
|
||||
assert!((rendered_mean(&mut x.s, id) - plain).abs() < 0.01);
|
||||
assert_eq!(x.s.media.denoise_salt(id, &x.s.catalog.photo(id).unwrap().develop), 0);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn what_was_made_is_found_again_by_a_new_session_and_a_new_model_starts_over() {
|
||||
let mut x = setup("again", true);
|
||||
let id = photo(&x.s);
|
||||
set_amount(&mut x.s, 100.0);
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
let made = x.runs.load(Ordering::Relaxed);
|
||||
x.s.close_library().unwrap();
|
||||
drop(x.s);
|
||||
|
||||
// reopened: ready without running the model again
|
||||
let mut s2 = Session::new().with_fs();
|
||||
s2.set_denoise_models_dir(Some(x.dir.join("denoise-models")));
|
||||
s2.denoise.loader = dim_loader(&x.runs);
|
||||
s2.open_library(x.dir.join("library"), false).unwrap();
|
||||
s2.denoise_pump(crate::denoise::Pace::Full);
|
||||
assert_eq!(s2.denoise_photo_state(id), PhotoState::Ready);
|
||||
assert_eq!(x.runs.load(Ordering::Relaxed), made, "nothing was made again");
|
||||
|
||||
// another model's pictures are different pictures
|
||||
let onnx = model_files(&x.dir, "other");
|
||||
s2.execute("denoise.models.install", &json!({"path": onnx.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
s2.denoise_pump(crate::denoise::Pace::Pause);
|
||||
assert!(
|
||||
matches!(s2.denoise_photo_state(id), PhotoState::Queued { .. }),
|
||||
"the first model's picture is not this model's: it is to be made: {:?}",
|
||||
s2.denoise_photo_state(id)
|
||||
);
|
||||
// and choosing the first again finds its picture where it was
|
||||
s2.execute("denoise.models.select", &json!({"id": "stand-in"})).unwrap();
|
||||
s2.denoise_pump(crate::denoise::Pace::Pause);
|
||||
assert_eq!(s2.denoise_photo_state(id), PhotoState::Ready);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_export_makes_its_picture_when_the_queue_has_not_and_never_comes_out_without_it() {
|
||||
let mut x = setup("export", true);
|
||||
let id = photo(&x.s);
|
||||
let o = ExportOptions::from_json(&json!({"format": "png", "width": 160, "height": 160}));
|
||||
let decode = |bytes: &[u8]| {
|
||||
let d = lightcraft_codecs::decode(bytes, lightcraft_codecs::DecodeOptions::fit(160, 160)).unwrap();
|
||||
let w = d.to_working();
|
||||
w.data.iter().map(|p| f64::from(p[0] + p[1] + p[2])).sum::<f64>() / (w.data.len() * 3) as f64
|
||||
};
|
||||
let plain = decode(&export_photo(&mut x.s, id, &o, 1).unwrap().bytes);
|
||||
assert_eq!(x.runs.load(Ordering::Relaxed), 0, "no amount, no denoise");
|
||||
set_amount(&mut x.s, 100.0);
|
||||
// no pump: the export makes the picture itself
|
||||
let denoised = decode(&export_photo(&mut x.s, id, &o, 1).unwrap().bytes);
|
||||
assert!(x.runs.load(Ordering::Relaxed) >= 4);
|
||||
assert!(denoised < plain * 0.9, "denoised {denoised} vs plain {plain}");
|
||||
let products = std::fs::read_dir(x.dir.join("library").join("denoise")).unwrap().flatten().count();
|
||||
assert_eq!(products, 1);
|
||||
// a second export finds it
|
||||
let runs = x.runs.load(Ordering::Relaxed);
|
||||
let again = decode(&export_photo(&mut x.s, id, &o, 1).unwrap().bytes);
|
||||
assert_eq!(x.runs.load(Ordering::Relaxed), runs);
|
||||
assert!((again - denoised).abs() < 1e-6);
|
||||
// the pump takes in what the export made
|
||||
x.s.denoise_pump(crate::denoise::Pace::Pause);
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready);
|
||||
|
||||
// no model: the export says so instead of leaving the denoise out
|
||||
x.s.execute("denoise.models.select", &json!({"id": null})).unwrap();
|
||||
let e = export_photo(&mut x.s, id, &o, 1).err().unwrap();
|
||||
assert!(e.contains("no denoise model is chosen") && e.contains("ramp.dng"), "{e}");
|
||||
set_amount(&mut x.s, 0.0);
|
||||
assert!(export_photo(&mut x.s, id, &o, 1).is_ok(), "without an amount nothing is asked of the model");
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_model_that_fails_fails_the_export_and_is_not_retried_by_the_queue() {
|
||||
let mut x = setup("fails", true);
|
||||
let id = photo(&x.s);
|
||||
x.s.denoise.loader = Arc::new(|_, _| Err("the model went away".into()));
|
||||
// Simulate loss of the in-memory model; successful installation now retains its tested instance.
|
||||
*x.s.denoise.active.as_ref().unwrap().model.lock().unwrap() = None;
|
||||
set_amount(&mut x.s, 80.0);
|
||||
let o = ExportOptions::from_json(&json!({"format": "png", "width": 100, "height": 100}));
|
||||
let e = export_photo(&mut x.s, id, &o, 1).err().unwrap();
|
||||
assert!(e.contains("ramp.dng") && e.contains("the model went away"), "{e}");
|
||||
// the background job fails too, says why once and does not try again every frame
|
||||
for _ in 0..200 {
|
||||
x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap();
|
||||
if matches!(x.s.denoise_photo_state(id), PhotoState::Failed { .. }) {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
match x.s.denoise_photo_state(id) {
|
||||
PhotoState::Failed { why, unsupported } => assert!(why.contains("the model went away") && !unsupported, "{why}"),
|
||||
other => panic!("{other:?}"),
|
||||
}
|
||||
let r = x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap();
|
||||
assert_eq!((r["queued"].clone(), r["running"].clone()), (json!(0), Value::Null), "{r}");
|
||||
// asking again by hand does retry
|
||||
x.s.denoise.loader = dim_loader(&x.runs);
|
||||
let r = x.s.execute("denoise.queue", &json!({"retry": true})).unwrap();
|
||||
assert_eq!(r["added"], 1, "{r}");
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_corrupt_cache_with_a_current_header_never_silently_exports_the_plain_photo() {
|
||||
let mut x = setup("corrupt-export", true);
|
||||
let id = photo(&x.s);
|
||||
set_amount(&mut x.s, 100.0);
|
||||
let options = ExportOptions::from_json(&json!({"format": "png", "width": 160, "height": 160}));
|
||||
export_photo(&mut x.s, id, &options, 1).unwrap();
|
||||
x.s.denoise_pump(crate::denoise::Pace::Pause);
|
||||
let spec = x.s.media.denoise.spec(id).unwrap().clone();
|
||||
let mut bytes = std::fs::read(&spec.product).unwrap();
|
||||
for byte in bytes.iter_mut().rev().take(64) {
|
||||
*byte ^= 0xff;
|
||||
}
|
||||
std::fs::write(&spec.product, bytes).unwrap();
|
||||
assert!(lightcraft_denoise::product::is_current(&spec.product, &spec.key), "header/length still look current");
|
||||
x.s.media.forget(id);
|
||||
let error = export_photo(&mut x.s, id, &options, 1).err().unwrap();
|
||||
assert!(error.contains("AI Denoise picture could not be read"), "{error}");
|
||||
let _ = std::fs::remove_dir_all(x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_raw_photos_on_disk_are_denoised_and_nothing_is_queued_without_a_model() {
|
||||
let dir = temp_dir("eligible");
|
||||
let mut demo = Session::with_demo();
|
||||
demo.set_denoise_models_dir(Some(dir.join("models")));
|
||||
let id = demo.active().unwrap();
|
||||
assert_eq!(demo.denoise_photo_state(id), PhotoState::NotApplicable, "a generated photo has no sensor data");
|
||||
let e = demo.execute("denoise.queue", &json!({})).unwrap_err().to_string();
|
||||
assert!(e.contains("no denoise model is chosen"), "{e}");
|
||||
// no models folder at all (the web build): the commands say so, nothing panics
|
||||
let mut web = Session::with_demo();
|
||||
for cmd in ["denoise.models.install", "denoise.models.select", "denoise.models.remove", "denoise.settings"] {
|
||||
assert!(web.execute(cmd, &json!({"path": "x", "acknowledged": true, "id": "x"})).is_err(), "{cmd}");
|
||||
}
|
||||
assert_eq!(web.execute("denoise.pump", &json!({})).unwrap()["active"], false);
|
||||
assert_eq!(web.execute("denoise.status", &json!({})).unwrap()["enabled"], false);
|
||||
// no library open: nowhere to keep pictures, and the queue says so
|
||||
let x = setup("eligible2", true);
|
||||
let mut bare = Session::new().with_fs();
|
||||
bare.set_denoise_models_dir(Some(x.dir.join("denoise-models")));
|
||||
assert!(bare.denoise_products_dir().is_none());
|
||||
let e = bare.execute("denoise.queue", &json!({})).unwrap_err().to_string();
|
||||
assert!(e.contains("open a library on disk first"), "{e}");
|
||||
assert_eq!(bare.execute("denoise.pump", &json!({})).unwrap()["queued"], 0);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn files_that_cannot_be_denoised_fail_with_a_reason_and_never_a_panic() {
|
||||
let dir = temp_dir("makefails");
|
||||
let spec = |path: &Path| crate::denoise::JobSpec {
|
||||
path: path.to_string_lossy().into_owned(),
|
||||
read: None,
|
||||
product: dir.join("p.lcdn"),
|
||||
key: "k".into(),
|
||||
manifest: manifest("stand-in"),
|
||||
onnx: dir.join("model.onnx"),
|
||||
model: Default::default(),
|
||||
loader: dim_loader(&Arc::new(AtomicUsize::new(0))),
|
||||
parallel: 2,
|
||||
run_on: RunOn::Auto,
|
||||
};
|
||||
use crate::denoise::{MakeError, make_product};
|
||||
// not a file
|
||||
assert!(matches!(make_product(&spec(&dir.join("missing.dng")), None), Err(MakeError::Failed(_))));
|
||||
// not a raw
|
||||
let junk = dir.join("junk.dng");
|
||||
std::fs::write(&junk, b"this is not a raw file at all").unwrap();
|
||||
assert!(matches!(make_product(&spec(&junk), None), Err(MakeError::Unsupported(_))));
|
||||
// a raw cut short
|
||||
let whole = ramp_dng(192, 160, "RGGB");
|
||||
let cut = dir.join("cut.dng");
|
||||
std::fs::write(&cut, &whole[..whole.len() / 3]).unwrap();
|
||||
assert!(make_product(&spec(&cut), None).is_err());
|
||||
// a model that gives the wrong number of values is an error from the model, not a crash
|
||||
struct Wrong;
|
||||
impl TileRunner for Wrong {
|
||||
fn run(&self, _: &[f32]) -> Result<Vec<f32>, lightcraft_denoise::Error> {
|
||||
Ok(vec![0.0; 7])
|
||||
}
|
||||
}
|
||||
impl Model for Wrong {
|
||||
fn runner(&self, _: RunOn, threads: usize) -> (&dyn TileRunner, usize) {
|
||||
(self, threads)
|
||||
}
|
||||
fn self_test(&self, _: RunOn) -> Result<Value, String> {
|
||||
Err("no".into())
|
||||
}
|
||||
}
|
||||
let good = dir.join("good.dng");
|
||||
std::fs::write(&good, &whole).unwrap();
|
||||
let mut s = spec(&good);
|
||||
s.loader = Arc::new(|_, _| Ok(Arc::new(Wrong) as Arc<dyn Model>));
|
||||
match make_product(&s, None) {
|
||||
Err(MakeError::Failed(why)) => assert!(why.contains("7 numbers"), "{why}"),
|
||||
other => panic!("{other:?}"),
|
||||
}
|
||||
assert!(!dir.join("p.lcdn").exists(), "a failed job leaves no product");
|
||||
// a job told to stop stops before doing anything
|
||||
let progress = crate::denoise::Progress::default();
|
||||
progress.cancel_now();
|
||||
assert_eq!(make_product(&spec(&good), Some(&progress)), Err(MakeError::Cancelled));
|
||||
// a Bayer layout of the other kinds is denoised too, and a product that exists is not made again
|
||||
for cfa in ["BGGR", "GRBG", "GBRG"] {
|
||||
let f = dir.join(format!("{cfa}.dng"));
|
||||
std::fs::write(&f, ramp_dng(192, 160, cfa)).unwrap();
|
||||
let mut s = spec(&f);
|
||||
s.product = dir.join(format!("{cfa}.lcdn"));
|
||||
assert_eq!(make_product(&s, None), Ok(true), "{cfa}");
|
||||
assert_eq!(make_product(&s, None), Ok(false), "{cfa} is already there");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_threads_making_the_same_picture_make_it_once() {
|
||||
let dir = temp_dir("twice");
|
||||
let f = dir.join("a.dng");
|
||||
std::fs::write(&f, ramp_dng(192, 160, "RGGB")).unwrap();
|
||||
let runs = Arc::new(AtomicUsize::new(0));
|
||||
let spec = crate::denoise::JobSpec {
|
||||
path: f.to_string_lossy().into_owned(),
|
||||
read: None,
|
||||
product: dir.join("p.lcdn"),
|
||||
key: "k".into(),
|
||||
manifest: manifest("stand-in"),
|
||||
onnx: dir.join("model.onnx"),
|
||||
model: Default::default(),
|
||||
loader: dim_loader(&runs),
|
||||
parallel: 2,
|
||||
run_on: RunOn::Auto,
|
||||
};
|
||||
let made: Vec<bool> = std::thread::scope(|sc| {
|
||||
let hs: Vec<_> = (0..4).map(|_| sc.spawn(|| crate::denoise::make_product(&spec, None).unwrap())).collect();
|
||||
hs.into_iter().map(|h| h.join().unwrap()).collect()
|
||||
});
|
||||
assert_eq!(made.iter().filter(|m| **m).count(), 1, "{made:?}");
|
||||
assert_eq!(runs.load(Ordering::Relaxed), 4, "one photo's four tiles, once");
|
||||
assert!(lightcraft_denoise::product::is_current(&spec.product, "k"));
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_cache_is_kept_within_its_limit_and_what_photos_use_goes_last() {
|
||||
let mut x = setup("evict", true);
|
||||
let id = photo(&x.s);
|
||||
set_amount(&mut x.s, 100.0);
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
let used = x.dir.join("library").join("denoise");
|
||||
// two older pictures nobody uses
|
||||
for (i, name) in ["aaaa", "bbbb"].iter().enumerate() {
|
||||
let f = used.join(format!("{name}.lcdn"));
|
||||
std::fs::write(&f, vec![7u8; 1000]).unwrap();
|
||||
let t = std::time::SystemTime::now() - std::time::Duration::from_secs(1000 - i as u64);
|
||||
std::fs::File::options().write(true).open(&f).unwrap().set_modified(t).unwrap();
|
||||
}
|
||||
let size_of = |p: &Path| std::fs::metadata(p).unwrap().len();
|
||||
let mine: PathBuf = std::fs::read_dir(&used).unwrap().flatten().map(|e| e.path()).find(|p| p.file_stem().is_some_and(|s| s.len() == 32)).unwrap();
|
||||
// room for the photo's own picture and one stray: the oldest stray goes, then the other
|
||||
x.s.denoise_evict_to(size_of(&mine) + 1500, None);
|
||||
assert!(!used.join("aaaa.lcdn").exists() && used.join("bbbb.lcdn").exists() && mine.exists());
|
||||
x.s.denoise_evict_to(size_of(&mine) + 10, None);
|
||||
assert!(!used.join("bbbb.lcdn").exists() && mine.exists());
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready);
|
||||
// too small for even that: the photo's own picture is the last thing to go, and the photo knows
|
||||
x.s.denoise_evict_to(10, None);
|
||||
assert!(!mine.exists());
|
||||
assert_ne!(x.s.denoise_photo_state(id), PhotoState::Ready);
|
||||
// …and it is made again, because the photo still has its amount
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
assert!(mine.exists());
|
||||
// clearing deletes it all
|
||||
let r = x.s.execute("denoise.clear", &json!({})).unwrap();
|
||||
assert_eq!(r["deleted"], 1, "{r}");
|
||||
assert_eq!(std::fs::read_dir(&used).unwrap().flatten().count(), 0);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_are_checked_saved_and_a_paused_pump_starts_nothing() {
|
||||
let mut x = setup("settings", true);
|
||||
let r = x.s.execute("denoise.settings", &json!({"auto": false, "cacheGb": 5, "threads": 3})).unwrap();
|
||||
assert_eq!((r["auto"].clone(), r["cacheGb"].clone(), r["threads"].clone()), (json!(false), json!(5), json!(3)), "{r}");
|
||||
for bad in [
|
||||
json!({"cacheGb": 0}),
|
||||
json!({"cacheGb": "big"}),
|
||||
json!({"threads": 0}),
|
||||
json!({"threads": 1000}),
|
||||
json!({"auto": "yes"}),
|
||||
json!({"runOn": "fastest"}),
|
||||
json!({"runOn": true}),
|
||||
] {
|
||||
assert!(x.s.execute("denoise.settings", &bad).is_err(), "{bad}");
|
||||
}
|
||||
assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["cacheGb"], 5, "a refused change changes nothing");
|
||||
// the faster of the card and the processor unless the user says otherwise, and the interface can read the choice back
|
||||
assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["runOn"], "auto");
|
||||
assert_eq!(x.s.execute("denoise.settings", &json!({"runOn": "cpu"})).unwrap()["runOn"], "cpu");
|
||||
assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["runOn"], "cpu");
|
||||
let status = x.s.execute("denoise.status", &json!({})).unwrap();
|
||||
assert_eq!((status["runOn"].clone(), status["device"]["kind"].clone()), (json!("cpu"), json!("none")), "no model is loaded yet: {status}");
|
||||
// the setting before `runOn` still reads as what it meant
|
||||
let models = x.s.denoise.models_dir.clone().unwrap();
|
||||
let mut st = crate::denoise::read_settings(&models);
|
||||
(st.run_on, st.gpu) = (None, Some(false));
|
||||
crate::denoise::write_settings(&models, &st).unwrap();
|
||||
assert_eq!(x.s.execute("denoise.models.list", &json!({})).unwrap()["runOn"], "cpu");
|
||||
// choosing where it runs lets the card be tried again after a set-up that closed LightCraft
|
||||
let onnx = crate::denoise::installed_models(&models).first().unwrap().onnx.clone();
|
||||
let marker = onnx.with_file_name(crate::denoise::GPU_SETUP_MARKER);
|
||||
std::fs::write(&marker, "setting up").unwrap();
|
||||
assert_eq!(x.s.execute("denoise.settings", &json!({"runOn": "auto"})).unwrap()["runOn"], "auto");
|
||||
assert!(!marker.exists(), "the crash is forgotten");
|
||||
assert_eq!(crate::denoise::read_settings(&models).gpu, None, "the old setting is not kept beside the new one");
|
||||
// with auto off nothing is started by itself, and the photo is made when asked
|
||||
let id = photo(&x.s);
|
||||
set_amount(&mut x.s, 100.0);
|
||||
for _ in 0..3 {
|
||||
x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap();
|
||||
}
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle);
|
||||
x.s.execute("denoise.settings", &json!({"auto": true})).unwrap();
|
||||
// paused: it is queued but nothing starts
|
||||
for _ in 0..3 {
|
||||
let r = x.s.execute("denoise.pump", &json!({"pace": "pause"})).unwrap();
|
||||
assert_eq!(r["running"], Value::Null);
|
||||
}
|
||||
assert!(matches!(x.s.denoise_photo_state(id), PhotoState::Queued { .. }), "{:?}", x.s.denoise_photo_state(id));
|
||||
// …and cancelling forgets it
|
||||
assert_eq!(x.s.execute("denoise.cancel", &json!({})).unwrap()["cancelled"], 1);
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle);
|
||||
// when the pace allows, it is made
|
||||
x.s.execute("denoise.queue", &json!({"scope": "withAmount"})).unwrap();
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_photo_whose_file_changed_has_no_picture_until_it_is_made_again() {
|
||||
let mut x = setup("changed", true);
|
||||
let id = photo(&x.s);
|
||||
set_amount(&mut x.s, 100.0);
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
// the photo's content is something else now (a replaced file): its old picture is not its picture
|
||||
let p = x.s.catalog.photo(id).unwrap().clone();
|
||||
x.s.commit(
|
||||
"test",
|
||||
lightcraft_catalog::Op::SetContent {
|
||||
id,
|
||||
width: p.width,
|
||||
height: p.height,
|
||||
file_size: p.file_size + 1,
|
||||
content_hash: Some("ffffffffffffffffffffffffffffffff".into()),
|
||||
preview_only: None,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
x.s.denoise_pump(crate::denoise::Pace::Pause);
|
||||
assert_ne!(x.s.denoise_photo_state(id), PhotoState::Ready);
|
||||
// the old picture is not used by a render either
|
||||
assert_eq!(x.s.media.denoise_salt(id, &x.s.catalog.photo(id).unwrap().develop), 0);
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_photo_waits_for_its_turn_unless_pictures_are_not_made_on_their_own_or_have_nowhere_to_go() {
|
||||
let mut x = setup("waits", true);
|
||||
let id = photo(&x.s);
|
||||
set_amount(&mut x.s, 80.0);
|
||||
// automatic off: the open photo is left alone until it is asked for
|
||||
x.s.execute("denoise.settings", &json!({"auto": false})).unwrap();
|
||||
assert!(!x.s.denoise_auto());
|
||||
x.s.execute("denoise.pump", &json!({"pace": "full"})).unwrap();
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Idle);
|
||||
assert!(!x.s.denoise_busy());
|
||||
x.s.execute("denoise.queue", &json!({"ids": [id.0]})).unwrap();
|
||||
assert!(x.s.denoise_busy(), "queued with a model: the headless runs wait for it");
|
||||
pump_until_ready(&mut x.s, 1);
|
||||
assert_eq!(x.s.denoise_photo_state(id), PhotoState::Ready);
|
||||
assert!(!x.s.denoise_busy());
|
||||
// a session whose library is not on disk has nowhere to keep it: the photo says so instead of waiting for ever
|
||||
x.s.denoise.settings.auto = Some(true);
|
||||
if let Some(l) = x.s.library.as_mut() {
|
||||
l.on_disk = false;
|
||||
}
|
||||
match x.s.denoise_photo_state(id) {
|
||||
PhotoState::Failed { why, unsupported } => assert!(unsupported && why.contains("library folder"), "{why}"),
|
||||
other => panic!("{other:?}"),
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&x.dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_denoise_switch_validates_arguments_and_accepts_explicit_targets_without_selection() {
|
||||
let mut x = setup("toggle", false);
|
||||
let id = photo(&x.s);
|
||||
for params in [json!({"id": "bad"}), json!({"id": u64::MAX}), json!({"id": id.0, "enabled": "yes"})] {
|
||||
assert!(x.s.execute("denoise.toggle", ¶ms).is_err());
|
||||
}
|
||||
assert!(x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).is_err());
|
||||
assert_eq!(x.s.develop_of(id).unwrap().enhance.denoise, 0.0);
|
||||
let model = model_files(&x.dir, "toggle-model");
|
||||
x.s.execute("denoise.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
x.s.selection = Default::default();
|
||||
let undo = x.s.undo.len();
|
||||
let r = x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap();
|
||||
assert_eq!(r["amount"], 50.0);
|
||||
assert_eq!(x.s.undo.len(), undo + 1);
|
||||
x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": false})).unwrap();
|
||||
assert_eq!(x.s.develop_of(id).unwrap().denoise_amount(), 0.0);
|
||||
// the switch is its own setting: flipping it keeps the Amount (so it compares with and without), and an Amount of 0
|
||||
// does not switch it off
|
||||
x.s.execute("library.select", &json!({"ids": [id.0]})).unwrap();
|
||||
x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap();
|
||||
x.s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 80.0})).unwrap();
|
||||
let off = x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": false})).unwrap();
|
||||
assert_eq!(off["amount"], 80.0);
|
||||
let d = x.s.develop_of(id).unwrap();
|
||||
assert!(!d.enhance.denoise_enabled() && d.denoise_amount() == 0.0);
|
||||
assert_eq!(x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap()["amount"], 80.0);
|
||||
assert!((x.s.develop_of(id).unwrap().denoise_amount() - 0.8).abs() < 1e-6);
|
||||
x.s.execute("develop.set", &json!({"control": "enhance.denoise", "value": 0.0})).unwrap();
|
||||
assert!(x.s.develop_of(id).unwrap().enhance.denoise_enabled(), "sliding to 0 leaves it on");
|
||||
assert_eq!(x.s.execute("denoise.toggle", &json!({"id": id.0, "enabled": true})).unwrap()["amount"], 0.0);
|
||||
// with no `enabled`, the toggle flips what the switch says
|
||||
assert_eq!(x.s.execute("denoise.toggle", &json!({"id": id.0})).unwrap()["enabled"], false);
|
||||
drop(x.s);
|
||||
let _ = std::fs::remove_dir_all(x.dir);
|
||||
}
|
||||
@@ -0,0 +1,586 @@
|
||||
//! 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::tiny_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"], cfg!(not(target_arch = "wasm32")));
|
||||
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");
|
||||
// speed is a ratio to a ResNet-100 model, in words, for models that are not installed yet too; never milliseconds
|
||||
assert_eq!(find(&l, "sface-2021dec")["speedText"], "6.9× faster than a ResNet-100 model");
|
||||
assert_eq!(find(&l, "auraface-v1")["speedText"], "Same speed as a ResNet-100 model");
|
||||
assert_eq!(find(&l, "yunet-2023mar")["speedText"], Value::Null, "nobody timed the detector against it");
|
||||
// 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);
|
||||
}
|
||||
|
||||
/// A demo session with the real YuNet installed, for the tests that need it: `LC_YUNET_MODEL=<the .onnx file>` (what
|
||||
/// Settings > Faces downloads; nothing is committed). Without it those tests say so and pass.
|
||||
fn demo_with_yunet() -> Option<Session> {
|
||||
let Some(file) = std::env::var_os("LC_YUNET_MODEL") else {
|
||||
eprintln!("LC_YUNET_MODEL is not set: skipping");
|
||||
return None;
|
||||
};
|
||||
let dir = temp("yunet");
|
||||
let mut s = Session::with_demo();
|
||||
s.face_models_dir = Some(dir.join("models"));
|
||||
let r = s.execute("faces.models.install", &json!({"path": file.to_string_lossy(), "acknowledged": true}));
|
||||
assert!(r.is_ok(), "{r:?}");
|
||||
Some(s)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn installing_the_detector_does_not_choose_it_as_the_recogniser_or_switch_recognition_on() {
|
||||
let Some(mut s) = demo_with_yunet() else { return };
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!(find(&l, "yunet-2023mar")["installed"], true);
|
||||
assert_eq!(l["enabled"], false, "{l}");
|
||||
assert!(l["embedder"].is_null(), "{l}");
|
||||
assert!(s.execute("faces.models.test", &json!({"id": "yunet-2023mar"})).is_err(), "a detector has no recognition self-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_without_the_model_says_where_to_get_it() {
|
||||
let d = temp("nodetector");
|
||||
let mut s = Session::with_demo();
|
||||
// no folder for models at all, then a folder that has none yet
|
||||
assert!(s.execute("faces.detect", &json!({})).is_err());
|
||||
s.face_models_dir = Some(d.join("models"));
|
||||
let e = s.execute("faces.detect", &json!({})).unwrap_err().to_string();
|
||||
assert!(e.contains("Settings > Faces"), "{e}");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_detector_and_the_recognisers_can_be_downloaded_and_the_download_is_refused_otherwise() {
|
||||
let d = temp("dl-refusals");
|
||||
let mut s = session(&d);
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
for (id, host) in [("yunet-2023mar", "github.com"), ("sface-2021dec", "github.com"), ("auraface-v1", "huggingface.co")] {
|
||||
assert_eq!(find(&l, id)["downloadHost"], host, "{id}");
|
||||
}
|
||||
// without the user's acceptance, for an id nobody knows, and without a folder: nothing starts
|
||||
assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar"})).is_err());
|
||||
assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": false})).is_err());
|
||||
for id in ["nope", "", "../yunet-2023mar"] {
|
||||
assert!(s.execute("faces.models.download", &json!({"id": id, "acknowledged": true})).is_err(), "{id}");
|
||||
}
|
||||
assert!(Session::new().execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).is_err());
|
||||
assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([]));
|
||||
assert!(!d.join("models").join(".downloads").exists(), "a refused download touches nothing");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_download_that_has_arrived_is_installed_by_itself_and_its_staged_file_goes() {
|
||||
let d = temp("dl-arrived");
|
||||
let mut s = session(&d);
|
||||
let staging = d.join("models").join(".downloads");
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let staged = staging.join("model.onnx");
|
||||
std::fs::write(&staged, lightcraft_faces::synthetic::tiny_embedder_model(512)).unwrap();
|
||||
let sha = lightcraft_faces::hash::sha256_file(&staged).unwrap();
|
||||
s.face_downloads.arrived("my-model", staged.clone(), &sha);
|
||||
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
let row = &r["downloads"][0];
|
||||
assert_eq!((row["id"].as_str(), row["state"].as_str()), (Some("my-model"), Some("installed")), "{r}");
|
||||
assert!(!staged.exists(), "the staged file was moved into the model's folder");
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!(l["models"].as_array().unwrap().iter().filter(|m| m["installed"] == true).count(), 1);
|
||||
// announced once, cleared on request
|
||||
assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "my-model"})).unwrap()["discarded"], true);
|
||||
assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([]));
|
||||
assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "my-model"})).unwrap()["discarded"], false);
|
||||
|
||||
// a file that cannot be installed is thrown away with a reason
|
||||
std::fs::write(&staged, b"this is not a model").unwrap();
|
||||
let sha = lightcraft_faces::hash::sha256_file(&staged).unwrap();
|
||||
s.face_downloads.arrived("bad", staged.clone(), &sha);
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
assert_eq!(r["downloads"][0]["state"], "failed", "{r}");
|
||||
assert!(!staged.exists());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
/// `faces.detect` on photos without faces: it reports nothing, and a new run replaces earlier detections but
|
||||
/// never regions that came from XMP or were named; one undo step restores everything.
|
||||
#[test]
|
||||
fn detect_replaces_only_earlier_detections_and_is_one_undo_step() {
|
||||
use lightcraft_catalog::Op;
|
||||
use lightcraft_meta::{Region, RegionKind};
|
||||
let region = |name: Option<&str>, description: Option<&str>, x: f64| Region {
|
||||
rect: lightcraft_geom::Rect { x0: x, y0: 0.2, x1: x + 0.2, y1: 0.5 },
|
||||
kind: RegionKind::Face,
|
||||
name: name.map(str::to_string),
|
||||
description: description.map(str::to_string),
|
||||
};
|
||||
let Some(mut s) = demo_with_yunet() else { return };
|
||||
let id = s.active().unwrap();
|
||||
let mut meta = s.catalog.photo(id).unwrap().meta.clone();
|
||||
meta.regions =
|
||||
vec![region(Some("Jane Doe"), None, 0.1), region(None, Some("Detected by YuNet 2023mar"), 0.5), region(None, Some("Drawn by hand"), 0.7)];
|
||||
s.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap();
|
||||
let undo_before = s.undo.len();
|
||||
|
||||
// a dry run reports and changes nothing
|
||||
let r = s.execute("faces.detect", &json!({"apply": false})).unwrap();
|
||||
assert_eq!(r["applied"], false);
|
||||
assert_eq!(s.catalog.photo(id).unwrap().meta.regions.len(), 3);
|
||||
assert_eq!(s.undo.len(), undo_before);
|
||||
|
||||
// a real run (the default) drops the earlier detection (the demo has no faces) and keeps the others
|
||||
let r = s.execute("faces.detect", &json!({})).unwrap();
|
||||
assert_eq!(r["applied"], true);
|
||||
assert_eq!(r["detector"], "YuNet 2023mar");
|
||||
let photo = r["photos"].as_array().unwrap().iter().find(|p| p["id"] == id.0).unwrap();
|
||||
assert_eq!(photo["faces"].as_array().unwrap().len(), 0, "the procedural demo photos have no faces");
|
||||
let names: Vec<_> = s.catalog.photo(id).unwrap().meta.regions.iter().map(|r| (r.name.clone(), r.description.clone())).collect();
|
||||
assert_eq!(names, vec![(Some("Jane Doe".to_string()), None), (None, Some("Drawn by hand".to_string()))]);
|
||||
assert_eq!(s.undo.len(), undo_before + 1, "one undo step for the whole run");
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert_eq!(s.catalog.photo(id).unwrap().meta.regions.len(), 3, "undo brings the detection back");
|
||||
|
||||
// odd parameters are tolerated, and nothing is detected without a selection
|
||||
for p in [json!({"score": 5}), json!({"score": -1}), json!({"nmsIou": 99}), json!({"maxFaces": 0}), json!({"score": "x"})] {
|
||||
assert!(s.execute("faces.detect", &p).is_ok(), "{p}");
|
||||
}
|
||||
let mut empty = Session::new();
|
||||
assert!(empty.execute("faces.detect", &json!({})).is_err());
|
||||
}
|
||||
|
||||
/// Opt-in, needs the internet: `cargo test -p lightcraft-engine real_download -- --ignored --nocapture`. Downloads YuNet
|
||||
/// from its pinned address the way the app does, waits for it to be installed, then finds faces with it.
|
||||
#[test]
|
||||
#[ignore = "downloads from github.com"]
|
||||
fn real_download_of_yunet_installs_it_and_detects() {
|
||||
let d = temp("real-download");
|
||||
let mut s = Session::with_demo();
|
||||
s.face_models_dir = Some(d.join("models"));
|
||||
let started = s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).unwrap();
|
||||
assert_eq!(started["from"], "github.com");
|
||||
let t0 = std::time::Instant::now();
|
||||
loop {
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
let row = r["downloads"][0].clone();
|
||||
match row["state"].as_str() {
|
||||
Some("installed") => break,
|
||||
Some("running") | Some("done") => {}
|
||||
_ => panic!("{r}"),
|
||||
}
|
||||
assert!(t0.elapsed() < std::time::Duration::from_secs(120), "{r}");
|
||||
std::thread::sleep(std::time::Duration::from_millis(100));
|
||||
}
|
||||
println!("downloaded and installed in {:?}", t0.elapsed());
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!(find(&l, "yunet-2023mar")["installed"], true);
|
||||
let r = s.execute("faces.detect", &json!({"apply": false})).unwrap();
|
||||
assert_eq!(r["detector"], "YuNet 2023mar");
|
||||
// asking again once installed is refused
|
||||
assert!(s.execute("faces.models.download", &json!({"id": "yunet-2023mar", "acknowledged": true})).is_err());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
/// Opt-in, needs the internet: the recognisers are big, so this only checks that bytes start to arrive (through the
|
||||
/// hosts' redirects) and that cancelling leaves nothing behind.
|
||||
#[test]
|
||||
#[ignore = "downloads from github.com and huggingface.co"]
|
||||
fn real_downloads_of_the_recognisers_start_and_can_be_cancelled() {
|
||||
let d = temp("real-start");
|
||||
let mut s = Session::with_demo();
|
||||
s.face_models_dir = Some(d.join("models"));
|
||||
for id in ["sface-2021dec", "auraface-v1"] {
|
||||
s.execute("faces.models.download", &json!({"id": id, "acknowledged": true})).unwrap();
|
||||
let t0 = std::time::Instant::now();
|
||||
loop {
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
let row = r["downloads"].as_array().unwrap().iter().find(|x| x["id"] == id).unwrap().clone();
|
||||
if row["state"] == "running" && row["bytes"].as_u64().unwrap_or(0) > 100_000 {
|
||||
println!("{id}: {} of {} bytes after {:?}", row["bytes"], row["total"], t0.elapsed());
|
||||
break;
|
||||
}
|
||||
assert!(row["state"] == "running", "{r}");
|
||||
assert!(t0.elapsed() < std::time::Duration::from_secs(60), "{r}");
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
}
|
||||
s.execute("faces.models.downloadCancel", &json!({"id": id})).unwrap();
|
||||
}
|
||||
drop(s);
|
||||
let staging = d.join("models").join(".downloads");
|
||||
let left: Vec<_> = std::fs::read_dir(&staging).map(|r| r.flatten().map(|e| e.file_name()).collect()).unwrap_or_default();
|
||||
println!("left in staging: {left:?}");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
/// With the recognition runtime: installing runs the model's self-test, a model that does not work is refused
|
||||
/// and leaves nothing behind, and `faces.models.test` runs it again.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
#[test]
|
||||
fn installing_runs_the_self_test_and_refuses_models_that_do_not_work() {
|
||||
let d = temp("selftest");
|
||||
let mut s = session(&d);
|
||||
// a working model passes; its record keeps the result
|
||||
let f = model_file(&d, "Works.onnx", 64);
|
||||
let r = s.execute("faces.models.install", &json!({"path": f, "acknowledged": true})).unwrap();
|
||||
let id = r["installed"]["id"].as_str().unwrap().to_string();
|
||||
assert_eq!(r["installed"]["accepted"]["selfTest"]["ok"], true, "{r}");
|
||||
assert_eq!(r["installed"]["accepted"]["selfTest"]["dimension"], 64);
|
||||
let t = s.execute("faces.models.test", &json!({"id": id})).unwrap();
|
||||
assert_eq!(t["ok"], true, "{t}");
|
||||
assert!(t["result"]["embedMs"].as_f64().unwrap() >= 0.0);
|
||||
assert!(s.execute("faces.models.test", &json!({"id": "nope"})).is_err());
|
||||
assert!(s.execute("faces.models.test", &json!({})).is_err());
|
||||
// a graph with no layers passes the shape check but cannot run: refused, and nothing is left on disk
|
||||
let broken = d.join("Broken.onnx");
|
||||
std::fs::write(&broken, lightcraft_faces::synthetic::embedder_model(64)).unwrap();
|
||||
let e = s.execute("faces.models.install", &json!({"path": broken.to_string_lossy(), "acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("could not be loaded") || e.contains("self-test"), "{e}");
|
||||
let leftovers: Vec<_> = std::fs::read_dir(d.join("models"))
|
||||
.unwrap()
|
||||
.flatten()
|
||||
.map(|e| e.file_name().to_string_lossy().into_owned())
|
||||
.filter(|n| n != "settings.json")
|
||||
.collect();
|
||||
assert!(leftovers.len() == 1 && leftovers[0].starts_with("custom-works-"), "only the working model remains: {leftovers:?}");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
#[test]
|
||||
fn failed_replacement_keeps_the_installed_model_and_records() {
|
||||
let d = temp("replacement");
|
||||
let models = d.join("models");
|
||||
std::fs::create_dir_all(&models).unwrap();
|
||||
let good = model_file(&d, "Good.onnx", 64);
|
||||
let write_catalog = |path: &std::path::Path| {
|
||||
let catalog = json!({"models": [{
|
||||
"id": "replacement", "name": "Replacement", "version": "1", "role": "embedder",
|
||||
"url": "https://models.example.org/weights.onnx",
|
||||
"sha256": lightcraft_faces::hash::sha256_file(path).unwrap(),
|
||||
"sizeBytes": std::fs::metadata(path).unwrap().len(),
|
||||
"licence": {"name": "MIT", "commercial": "yes"},
|
||||
"output": {"kind": "embedding", "dim": 64}
|
||||
}]});
|
||||
std::fs::write(models.join("catalog.json"), catalog.to_string()).unwrap();
|
||||
};
|
||||
write_catalog(std::path::Path::new(&good));
|
||||
let mut s = session(&d);
|
||||
s.execute("faces.models.install", &json!({"path": good, "acknowledged": true})).unwrap();
|
||||
let home = models.join("replacement");
|
||||
let files = ["model.onnx", "face-model.json", "installed.json"];
|
||||
let before: Vec<_> = files.iter().map(|name| std::fs::read(home.join(name)).unwrap()).collect();
|
||||
let settings = std::fs::read(models.join("settings.json")).unwrap();
|
||||
let bad = d.join("Bad.onnx");
|
||||
std::fs::write(&bad, lightcraft_faces::synthetic::embedder_model(64)).unwrap();
|
||||
write_catalog(&bad);
|
||||
let mut s = session(&d);
|
||||
assert!(s.execute("faces.models.install", &json!({"path": bad, "acknowledged": true})).is_err());
|
||||
for (name, bytes) in files.iter().zip(before) {
|
||||
assert_eq!(std::fs::read(home.join(name)).unwrap(), bytes, "{name} was replaced");
|
||||
}
|
||||
assert_eq!(std::fs::read(models.join("settings.json")).unwrap(), settings);
|
||||
assert!(!home.join("model.onnx.part").exists());
|
||||
let _ = std::fs::remove_dir_all(d);
|
||||
}
|
||||
|
||||
fn settings_of(d: &std::path::Path) -> Value {
|
||||
serde_json::from_slice(&std::fs::read(d.join("models").join("settings.json")).unwrap()).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn installing_a_model_chooses_it_and_switches_recognition_on() {
|
||||
let d = temp("activate");
|
||||
let mut s = session(&d);
|
||||
let a = s.execute("faces.models.install", &json!({"path": model_file(&d, "First.onnx", 64), "acknowledged": true})).unwrap();
|
||||
let a_id = a["installed"]["id"].as_str().unwrap().to_string();
|
||||
assert_eq!(a["active"], json!({"embedder": a_id, "enabled": true}));
|
||||
assert_eq!(a["installed"]["selected"], true);
|
||||
assert_eq!(settings_of(&d), json!({"enabled": true, "embedder": a_id}));
|
||||
// a second model takes over (the newest is the one the user just asked for); the first stays installed
|
||||
let b = s.execute("faces.models.install", &json!({"path": model_file(&d, "Second.onnx", 96), "acknowledged": true})).unwrap();
|
||||
let b_id = b["installed"]["id"].as_str().unwrap().to_string();
|
||||
assert_ne!(a_id, b_id);
|
||||
assert_eq!(b["active"]["embedder"], json!(b_id));
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!((find(&l, &a_id)["installed"].clone(), find(&l, &a_id)["selected"].clone()), (json!(true), json!(false)));
|
||||
assert_eq!(find(&l, &b_id)["selected"], true);
|
||||
// switching back is a choice, and a model added with `activate: false` leaves the choice alone
|
||||
s.execute("faces.models.select", &json!({"id": a_id})).unwrap();
|
||||
let c = s.execute("faces.models.install", &json!({"path": model_file(&d, "Third.onnx", 128), "acknowledged": true, "activate": false})).unwrap();
|
||||
let c_id = c["installed"]["id"].as_str().unwrap().to_string();
|
||||
assert_eq!(c["active"]["embedder"], json!(a_id));
|
||||
assert_eq!(c["installed"]["selected"], false);
|
||||
// a switched-off recognition is switched back on by the next install, not by a refused one
|
||||
s.execute("faces.enable", &json!({"enabled": false})).unwrap();
|
||||
assert!(s.execute("faces.models.install", &json!({"path": model_file(&d, "Nope.onnx", 64)})).is_err());
|
||||
assert_eq!(settings_of(&d)["enabled"], false);
|
||||
s.execute("faces.models.install", &json!({"path": model_file(&d, "Fourth.onnx", 64), "acknowledged": true})).unwrap();
|
||||
assert_eq!(settings_of(&d)["enabled"], true);
|
||||
// removing the model in use hands over to another installed one instead of leaving nothing chosen
|
||||
let in_use = settings_of(&d)["embedder"].as_str().unwrap().to_string();
|
||||
let r = s.execute("faces.models.remove", &json!({"id": in_use})).unwrap();
|
||||
let next = r["embedder"].as_str().unwrap().to_string();
|
||||
assert_ne!(next, in_use);
|
||||
assert!([&a_id, &b_id, &c_id].contains(&&next) || next.starts_with("custom-"), "{next}");
|
||||
assert_eq!(settings_of(&d)["embedder"], json!(next));
|
||||
// removing a model that is not in use changes nothing
|
||||
let other = [&a_id, &b_id, &c_id].into_iter().find(|i| **i != next && **i != in_use).unwrap().clone();
|
||||
s.execute("faces.models.remove", &json!({"id": other})).unwrap();
|
||||
assert_eq!(settings_of(&d)["embedder"], json!(next));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_download_that_arrives_is_installed_and_switched_on_by_itself() {
|
||||
use crate::face_download::{STAGING, State};
|
||||
let d = temp("finish");
|
||||
let mut s = session(&d);
|
||||
let staging = d.join("models").join(STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
// what the download thread leaves when it succeeds: a verified file in the staging folder and a finished job
|
||||
let staged = model_file(&staging, "arrived.onnx", 64);
|
||||
let sha = lightcraft_faces::hash::sha256_file(std::path::Path::new(&staged)).unwrap();
|
||||
s.face_downloads.set_outcome("fake-model", State::Done { path: staged.clone().into(), sha256: sha });
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
assert_eq!(r["downloads"][0]["id"], "fake-model");
|
||||
assert_eq!(r["downloads"][0]["state"], "installed", "{r}");
|
||||
let st = settings_of(&d);
|
||||
assert_eq!(st["enabled"], true);
|
||||
assert!(st["embedder"].as_str().unwrap().starts_with("custom-arrived-"), "{st}");
|
||||
assert!(!std::path::Path::new(&staged).exists(), "the staged file was moved into place");
|
||||
// the record stays until it is cleared, and is cleared once
|
||||
assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"][0]["state"], "installed");
|
||||
assert_eq!(s.execute("faces.models.downloadCancel", &json!({"id": "fake-model"})).unwrap()["discarded"], true);
|
||||
assert_eq!(s.execute("faces.models.downloads", &json!({})).unwrap()["downloads"], json!([]));
|
||||
|
||||
// a download that turns out not to be a usable model is thrown away, and the reason is shown
|
||||
let junk = staging.join("junk.onnx");
|
||||
std::fs::write(&junk, b"definitely not a model").unwrap();
|
||||
let sha = lightcraft_faces::hash::sha256_file(&junk).unwrap();
|
||||
s.face_downloads.set_outcome("fake-junk", State::Done { path: junk.clone(), sha256: sha });
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
assert_eq!(r["downloads"][0]["state"], "failed", "{r}");
|
||||
assert!(r["downloads"][0]["error"].as_str().unwrap().contains("cannot be used"), "{r}");
|
||||
assert!(!junk.exists());
|
||||
assert_eq!(settings_of(&d)["embedder"], st["embedder"], "a failed download leaves the choice alone");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_users_own_catalog_adds_models_with_the_same_download_and_install() {
|
||||
use crate::face_download::State;
|
||||
let d = temp("catalog");
|
||||
let mut s = session(&d);
|
||||
let models = d.join("models");
|
||||
std::fs::create_dir_all(&models).unwrap();
|
||||
// a model file, and a catalog that lists it by its hash (plus an entry that cannot be used)
|
||||
let file = model_file(&d, "weights.onnx", 64);
|
||||
let (sha, size) = (lightcraft_faces::hash::sha256_file(std::path::Path::new(&file)).unwrap(), std::fs::metadata(&file).unwrap().len());
|
||||
let catalog = json!({"models": [
|
||||
{"id": "my-research-model", "name": "My research model", "version": "1", "role": "embedder", "url": "https://models.example.org/w/weights.onnx",
|
||||
"sha256": sha, "sizeBytes": size, "licence": {"name": "Research only", "commercial": "no", "notice": "Not for commercial use."},
|
||||
"provenance": "Some dataset", "output": {"kind": "embedding", "dim": 64}, "thresholds": {"matchCosine": 0.42}},
|
||||
{"id": "broken", "name": "No hash", "version": "1", "role": "embedder", "url": "https://models.example.org/x.onnx", "output": {"kind": "embedding", "dim": 64}},
|
||||
]});
|
||||
std::fs::write(models.join("catalog.json"), catalog.to_string()).unwrap();
|
||||
// it is listed like a built-in model, with its own terms, a download address, and the problem with the other entry
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
let m = find(&l, "my-research-model");
|
||||
assert_eq!(
|
||||
(m["fromCatalog"].clone(), m["downloadHost"].clone(), m["licence"]["commercial"].clone()),
|
||||
(json!(true), json!("models.example.org"), json!("no"))
|
||||
);
|
||||
assert_eq!(m["installed"], false);
|
||||
assert_eq!(find(&l, "broken"), &Value::Null);
|
||||
assert_eq!(l["catalog"]["models"], 1);
|
||||
assert!(l["catalog"]["errors"][0].as_str().unwrap().contains("broken"), "{}", l["catalog"]);
|
||||
// nothing is fetched without the terms being accepted
|
||||
assert!(s.execute("faces.models.download", &json!({"id": "my-research-model"})).is_err());
|
||||
// a file that matches is recognised as that model (not as a draft), so its terms and settings are used
|
||||
let seen = s.execute("faces.models.inspect", &json!({"path": file})).unwrap();
|
||||
assert_eq!((seen["kind"].clone(), seen["model"]["id"].clone()), (json!("known"), json!("my-research-model")));
|
||||
// a finished download of it installs under the catalog's id, with the catalog's licence record
|
||||
let staging = models.join(crate::face_download::STAGING);
|
||||
std::fs::create_dir_all(&staging).unwrap();
|
||||
let staged = staging.join("weights.onnx");
|
||||
std::fs::copy(&file, &staged).unwrap();
|
||||
s.face_downloads.set_outcome("my-research-model", State::Done { path: staged, sha256: sha });
|
||||
let r = s.execute("faces.models.downloads", &json!({})).unwrap();
|
||||
assert_eq!(r["downloads"][0]["state"], "installed", "{r}");
|
||||
let st = settings_of(&d);
|
||||
assert_eq!((st["embedder"].clone(), st["enabled"].clone()), (json!("my-research-model"), json!(true)));
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert_eq!(
|
||||
(find(&l, "my-research-model")["installed"].clone(), find(&l, "my-research-model")["accepted"]["commercial"].clone()),
|
||||
(json!(true), json!("no"))
|
||||
);
|
||||
// a broken or hostile catalog never stops the list
|
||||
for bytes in [&b"{"[..], b"[]", &vec![b'x'; 400_000]] {
|
||||
std::fs::write(models.join("catalog.json"), bytes).unwrap();
|
||||
let l = s.execute("faces.models.list", &json!({})).unwrap();
|
||||
assert!(l["catalog"]["errors"].as_array().is_some_and(|e| !e.is_empty()), "{}", l["catalog"]);
|
||||
assert!(find(&l, "sface-2021dec")["id"] == "sface-2021dec");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
//! Recognition: indexing faces, suggesting names, naming, with a tiny working network standing in for a real
|
||||
//! model (its embedding depends on the picture's colours, which is enough to test the machinery; how well real
|
||||
//! models recognise people is measured separately, on real photos).
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use lightcraft_catalog::{Op, PhotoId};
|
||||
use lightcraft_geom::Rect;
|
||||
use lightcraft_meta::{Region, RegionKind};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::Session;
|
||||
|
||||
pub(crate) fn temp(name: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("lc-facerec-{name}-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
/// A demo session with a tiny model installed and chosen.
|
||||
pub(crate) fn setup(d: &std::path::Path, dim: u64) -> (Session, String) {
|
||||
let mut s = Session::with_demo();
|
||||
s.face_models_dir = Some(d.join("models"));
|
||||
let model = d.join(format!("tiny{dim}.onnx"));
|
||||
std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap();
|
||||
let r = s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
let id = r["installed"]["id"].as_str().unwrap().to_string();
|
||||
s.execute("faces.models.select", &json!({"id": id})).unwrap();
|
||||
(s, id)
|
||||
}
|
||||
|
||||
/// Installs the real YuNet from `LC_YUNET_MODEL=<the .onnx file>` (what Settings > Faces downloads; nothing is
|
||||
/// committed) and says whether it did: the parts of a test that need to find faces run only then.
|
||||
pub(crate) fn install_yunet(s: &mut Session) -> bool {
|
||||
let Some(file) = std::env::var_os("LC_YUNET_MODEL") else {
|
||||
eprintln!("LC_YUNET_MODEL is not set: skipping the face-finding part");
|
||||
return false;
|
||||
};
|
||||
let r = s.execute("faces.models.install", &json!({"path": file.to_string_lossy(), "acknowledged": true}));
|
||||
assert!(r.is_ok(), "{r:?}");
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) fn region(x: f64, name: Option<&str>) -> Region {
|
||||
Region { rect: Rect { x0: x, y0: 0.2, x1: x + 0.25, y1: 0.6 }, kind: RegionKind::Face, name: name.map(str::to_string), description: None }
|
||||
}
|
||||
|
||||
pub(crate) fn set_regions(s: &mut Session, id: PhotoId, regions: Vec<Region>) {
|
||||
let mut meta = s.catalog.photo(id).unwrap().meta.clone();
|
||||
meta.regions = regions;
|
||||
s.commit("setup", Op::SetMeta { id, meta: Box::new(meta) }).unwrap();
|
||||
}
|
||||
|
||||
pub(crate) fn two_photos(s: &Session) -> (PhotoId, PhotoId) {
|
||||
let mut ids = s.catalog.photos().map(|p| p.id).collect::<Vec<_>>();
|
||||
ids.sort();
|
||||
(ids[0], ids[1])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn indexing_embeds_each_face_once_and_reports_progress() {
|
||||
let d = temp("index");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
let (a, b) = two_photos(&s);
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]);
|
||||
set_regions(&mut s, b, vec![region(0.3, None)]);
|
||||
// only a report
|
||||
let r = s.execute("faces.index", &json!({"budgetMs": 0})).unwrap();
|
||||
assert_eq!((r["embedded"].clone(), r["photosDone"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(0), json!(2)), "{r}");
|
||||
// a generous budget does it all, the named photo first
|
||||
let r = s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap();
|
||||
assert_eq!((r["embedded"].clone(), r["pendingPhotos"].clone(), r["indexedFaces"].clone()), (json!(3), json!(0), json!(3)), "{r}");
|
||||
// nothing is embedded twice
|
||||
let r = s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap();
|
||||
assert_eq!((r["embedded"].clone(), r["photosDone"].clone()), (json!(0), json!(0)));
|
||||
// a small budget still makes progress (at least one photo per call), and `ids` go first
|
||||
let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap();
|
||||
set_regions(&mut s, c, vec![region(0.2, None)]);
|
||||
set_regions(&mut s, b, vec![region(0.3, None), region(0.6, None)]);
|
||||
let r = s.execute("faces.index", &json!({"budgetMs": 1, "ids": [c.0]})).unwrap();
|
||||
assert_eq!(r["photosDone"], 1, "{r}");
|
||||
assert_eq!(r["pendingPhotos"], 1, "{r}");
|
||||
assert!(s.faces.index.contains(c.0, ®ion(0.2, None).rect), "the photo asked for first was done first");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggestions_come_from_the_named_faces_and_are_never_applied() {
|
||||
let d = temp("suggest");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
let (a, b) = two_photos(&s);
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann"))]);
|
||||
set_regions(&mut s, b, vec![region(0.1, None), region(0.5, Some("Pet?")), region(0.7, None)]);
|
||||
// thresholds that accept anything: Ann is the only person with a named face except "Pet?"
|
||||
let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 60_000})).unwrap();
|
||||
assert_eq!(r["galleryFaces"], 2, "{r}");
|
||||
let faces = r["photos"][0]["faces"].as_array().unwrap();
|
||||
assert_eq!(faces.len(), 2, "only the unnamed faces get suggestions: {r}");
|
||||
assert!(faces.iter().all(|f| f["suggestion"]["name"].is_string()), "{r}");
|
||||
assert_eq!(faces[0]["index"], 0);
|
||||
assert_eq!(faces[1]["index"], 2);
|
||||
assert!(faces[0]["candidates"].as_array().unwrap().len() == 2);
|
||||
// an impossible threshold suggests nobody but still lists the candidates
|
||||
let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": 2.0, "budgetMs": 60_000})).unwrap();
|
||||
assert!(r["photos"][0]["faces"][0]["suggestion"].is_null());
|
||||
assert!(!r["photos"][0]["faces"][0]["candidates"].as_array().unwrap().is_empty());
|
||||
// nothing was named by suggesting
|
||||
assert!(s.catalog.photo(b).unwrap().meta.regions.iter().filter(|r| r.name.is_none()).count() == 2);
|
||||
// pets and other kinds are not people
|
||||
let mut meta = s.catalog.photo(b).unwrap().meta.clone();
|
||||
meta.regions[1].kind = RegionKind::Pet;
|
||||
s.commit("pet", Op::SetMeta { id: b, meta: Box::new(meta) }).unwrap();
|
||||
let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 60_000})).unwrap();
|
||||
assert_eq!(r["galleryFaces"], 1, "{r}");
|
||||
// odd parameters are tolerated
|
||||
for p in [json!({"threshold": "x"}), json!({"margin": null}), json!({"budgetMs": 0}), json!({"ids": []}), json!({"ids": [999999]})] {
|
||||
assert!(s.execute("faces.suggest", &p).is_ok(), "{p}");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn naming_a_face_is_undoable_and_makes_it_yours() {
|
||||
let d = temp("name");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
let (a, _) = two_photos(&s);
|
||||
let mut detected = region(0.1, None);
|
||||
detected.description = Some("Detected by YuNet 2023mar".into());
|
||||
set_regions(&mut s, a, vec![detected, region(0.5, None)]);
|
||||
let undo = s.undo.len();
|
||||
let r = s.execute("faces.setName", &json!({"id": a.0, "index": 0, "name": " Ann "})).unwrap();
|
||||
assert_eq!(r["name"], "Ann");
|
||||
let regions = s.catalog.photo(a).unwrap().meta.regions.clone();
|
||||
assert_eq!((regions[0].name.as_deref(), regions[0].description.as_deref()), (Some("Ann"), None), "a named detection loses its detected marker");
|
||||
assert_eq!(s.undo.len(), undo + 1);
|
||||
// a new detection run now leaves it alone
|
||||
if install_yunet(&mut s) {
|
||||
s.execute("faces.detect", &json!({"id": a.0})).unwrap();
|
||||
assert_eq!(s.catalog.photo(a).unwrap().meta.regions.iter().filter(|r| r.name.as_deref() == Some("Ann")).count(), 1);
|
||||
}
|
||||
// clearing, undo
|
||||
s.execute("faces.setName", &json!({"id": a.0, "index": 0, "name": null})).unwrap();
|
||||
assert_eq!(s.catalog.photo(a).unwrap().meta.regions[0].name, None);
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert_eq!(s.catalog.photo(a).unwrap().meta.regions[0].name.as_deref(), Some("Ann"));
|
||||
// bad input
|
||||
for p in [
|
||||
json!({"id": a.0}),
|
||||
json!({"id": a.0, "index": 9, "name": "x"}),
|
||||
json!({"id": a.0, "index": 0, "name": 5}),
|
||||
json!({"id": a.0, "index": 0, "name": "a\u{7}b"}),
|
||||
json!({"id": a.0, "index": 0, "name": "x".repeat(201)}),
|
||||
json!({"id": 999999, "index": 0, "name": "x"}),
|
||||
] {
|
||||
assert!(s.execute("faces.setName", &p).is_err(), "{p}");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embeddings_are_cached_beside_the_library_and_reset_for_another_model() {
|
||||
let d = temp("persist");
|
||||
let lib = d.join("lib");
|
||||
let models = d.join("models");
|
||||
let install = |s: &mut Session, dim: u64| -> String {
|
||||
let model = d.join(format!("tiny{dim}.onnx"));
|
||||
std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(dim)).unwrap();
|
||||
let r = s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
let id = r["installed"]["id"].as_str().unwrap().to_string();
|
||||
s.execute("faces.models.select", &json!({"id": id})).unwrap();
|
||||
id
|
||||
};
|
||||
let mut s = Session::new().with_fs();
|
||||
s.face_models_dir = Some(models.clone());
|
||||
s.open_library(&lib, true).unwrap();
|
||||
let (a, _) = two_photos(&s);
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]);
|
||||
install(&mut s, 64);
|
||||
s.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap();
|
||||
assert!(lib.join(format!("face-embeddings-{}.bin", s.faces.index.tag.split('@').next().unwrap())).is_file());
|
||||
s.close_library().unwrap();
|
||||
drop(s);
|
||||
|
||||
// a new session finds the embeddings on disk and has nothing left to do
|
||||
let mut s2 = Session::new().with_fs();
|
||||
s2.face_models_dir = Some(models);
|
||||
s2.open_library(&lib, true).unwrap();
|
||||
let r = s2.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap();
|
||||
assert_eq!((r["embedded"].clone(), r["indexedFaces"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(2), json!(0)), "{r}");
|
||||
// choosing another model starts over: its faces are different numbers, the old ones are not used
|
||||
install(&mut s2, 32);
|
||||
let r = s2.execute("faces.index", &json!({"budgetMs": 0})).unwrap();
|
||||
assert_eq!((r["indexedFaces"].clone(), r["pendingPhotos"].clone()), (json!(0), json!(1)), "{r}");
|
||||
let r = s2.execute("faces.index", &json!({"budgetMs": 60_000})).unwrap();
|
||||
assert_eq!(r["embedded"], 2);
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn without_a_chosen_model_recognition_says_so() {
|
||||
let d = temp("nomodel");
|
||||
let mut s = Session::with_demo();
|
||||
assert!(s.execute("faces.index", &json!({})).is_err(), "no models folder");
|
||||
s.face_models_dir = Some(d.join("models"));
|
||||
for cmd in ["faces.index", "faces.suggest"] {
|
||||
let e = s.execute(cmd, &json!({})).unwrap_err().to_string();
|
||||
assert!(e.contains("no recognition model is chosen"), "{cmd}: {e}");
|
||||
}
|
||||
let _: Value = json!(null);
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_background_pump_indexes_by_itself_once_recognition_is_on() {
|
||||
let d = temp("pump");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
let (a, b) = two_photos(&s);
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]);
|
||||
set_regions(&mut s, b, vec![region(0.3, None)]);
|
||||
// installing the model switched recognition on; switched off, the pump does nothing, and says so
|
||||
assert_eq!(s.execute("faces.enable", &json!({"enabled": false})).unwrap()["enabled"], false);
|
||||
assert_eq!(s.execute("faces.pump", &json!({})).unwrap()["active"], false);
|
||||
assert_eq!(s.faces.index.len(), 0);
|
||||
s.execute("faces.enable", &json!({"enabled": true})).unwrap();
|
||||
// `enabled` is looked up at most once a second, so the first pump may still see "off"
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
|
||||
let mut last = json!(null);
|
||||
while std::time::Instant::now() < deadline {
|
||||
last = s.execute("faces.pump", &json!({})).unwrap();
|
||||
if last["active"] == true && last["pendingPhotos"] == 0 && last["indexedFaces"] == 3 {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
}
|
||||
assert_eq!((last["active"].clone(), last["pendingPhotos"].clone(), last["indexedFaces"].clone()), (json!(true), json!(0), json!(3)), "{last}");
|
||||
// a new face is picked up (the catalog changed), and suggestions use what the pump embedded
|
||||
set_regions(&mut s, b, vec![region(0.3, None), region(0.6, None)]);
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
|
||||
while std::time::Instant::now() < deadline && s.faces.index.len() < 4 {
|
||||
s.execute("faces.pump", &json!({})).unwrap();
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
}
|
||||
assert_eq!(s.faces.index.len(), 4);
|
||||
let r = s.execute("faces.suggest", &json!({"ids": [b.0], "threshold": -1.0, "margin": -2.0, "budgetMs": 0})).unwrap();
|
||||
assert_eq!(r["photos"][0]["faces"].as_array().unwrap().len(), 2);
|
||||
assert!(r["photos"][0]["faces"][0]["suggestion"]["name"].is_string(), "{r}");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_background_pump_saves_what_it_embeds_for_the_next_launch() {
|
||||
let d = temp("pumpsave");
|
||||
let (lib, models) = (d.join("lib"), d.join("models"));
|
||||
let mut s = Session::new().with_fs();
|
||||
s.face_models_dir = Some(models);
|
||||
s.open_library(&lib, true).unwrap();
|
||||
let model = d.join("tiny64.onnx");
|
||||
std::fs::write(&model, lightcraft_faces::synthetic::tiny_embedder_model(64)).unwrap();
|
||||
// installing switches recognition on: nothing else is asked for
|
||||
s.execute("faces.models.install", &json!({"path": model.to_string_lossy(), "acknowledged": true})).unwrap();
|
||||
// without the detector (not bundled) photos without faces are left unsearched
|
||||
let searching = install_yunet(&mut s);
|
||||
let (a, b) = two_photos(&s);
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None)]);
|
||||
set_regions(&mut s, b, vec![region(0.3, None)]);
|
||||
let photos = s.catalog.photos().count();
|
||||
let searched_photos = if searching { photos - 2 } else { 0 };
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120);
|
||||
let mut last = json!(null);
|
||||
while std::time::Instant::now() < deadline {
|
||||
last = s.execute("faces.pump", &json!({})).unwrap();
|
||||
if last["pendingPhotos"] == 0 && s.faces.index.len() == 3 && s.faces.scanned.len() == searched_photos {
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
}
|
||||
assert_eq!((s.faces.index.len(), s.faces.scanned.len()), (3, searched_photos), "{last}");
|
||||
let (tag, path) = (s.faces.index.tag.clone(), s.faces.index.path.clone().unwrap());
|
||||
assert!(path.file_name().unwrap().to_string_lossy().starts_with("face-embeddings-custom-"), "one cache file per model: {}", path.display());
|
||||
// the first batch is written at once; whatever is left is written when the session ends, so a launch never redoes
|
||||
// work: neither the embeddings nor the list of photos already searched for faces
|
||||
s.close_library().unwrap();
|
||||
drop(s);
|
||||
let mut again = crate::faces_index::Index::default();
|
||||
again.reset(&tag, 64, Some(path));
|
||||
assert_eq!(again.len(), 3);
|
||||
let mut searched = crate::faces_index::Scanned::default();
|
||||
searched.reset(Some(lib.join("face-scanned.bin")));
|
||||
assert_eq!(searched.len(), searched_photos, "every photo without faces is remembered as searched");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn faces_person_lists_a_persons_faces_and_looks_for_more_only_while_recognition_runs() {
|
||||
let d = temp("person");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
s.execute("faces.index", &json!({"budgetMs": 0})).unwrap();
|
||||
let (a, b) = two_photos(&s);
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, Some("Bob"))]);
|
||||
set_regions(&mut s, b, vec![region(0.1, None), region(0.5, Some("ann"))]);
|
||||
let unit = |parts: &[(usize, f32)]| {
|
||||
let mut v = vec![0.0f32; 64];
|
||||
for (i, x) in parts {
|
||||
v[*i] = *x;
|
||||
}
|
||||
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
v.iter().map(|x| x / n).collect::<Vec<f32>>()
|
||||
};
|
||||
for (id, x, v) in [(a, 0.1, unit(&[(0, 1.0)])), (a, 0.5, unit(&[(1, 1.0)])), (b, 0.1, unit(&[(0, 0.9), (1, 0.1)])), (b, 0.5, unit(&[(0, 1.0)]))] {
|
||||
s.faces.index.insert(id.0, ®ion(x, None).rect, v);
|
||||
}
|
||||
// the name is matched without regard to case or surrounding spaces; every face is listed, wherever it is
|
||||
let r = s.execute("faces.person", &json!({"name": " ANN "})).unwrap();
|
||||
assert_eq!((r["name"].clone(), r["total"].clone()), (json!("Ann"), json!(2)));
|
||||
let mut at: Vec<(u64, u64)> =
|
||||
r["confirmed"].as_array().unwrap().iter().map(|f| (f["photo"].as_u64().unwrap(), f["index"].as_u64().unwrap())).collect();
|
||||
at.sort();
|
||||
let mut want = vec![(a.0, 0), (b.0, 1)];
|
||||
want.sort();
|
||||
assert_eq!(at, want);
|
||||
assert!(r["confirmed"].as_array().unwrap().iter().any(|f| f["rect"]["x0"] == 0.1 && f["rect"]["y1"] == 0.6));
|
||||
// recognition is on (installing the model switched it on): the unnamed face that looks like Ann is offered
|
||||
assert_eq!(r["ready"], true);
|
||||
let more = r["more"].as_array().unwrap();
|
||||
assert_eq!((more.len(), more[0]["photo"].clone(), more[0]["index"].clone()), (1, json!(b.0), json!(0)), "{r}");
|
||||
assert!(more[0]["score"].as_f64().unwrap() > 0.9);
|
||||
// confirming it moves it from one list to the other (and is one undo step)
|
||||
s.execute("faces.setName", &json!({"id": b.0, "index": 0, "name": "Ann"})).unwrap();
|
||||
let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap();
|
||||
assert_eq!((r["total"].clone(), r["more"].as_array().unwrap().len()), (json!(3), 0));
|
||||
// switched off, the faces are still listed but nothing is searched for
|
||||
s.execute("faces.enable", &json!({"enabled": false})).unwrap();
|
||||
let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap();
|
||||
assert_eq!((r["total"].clone(), r["ready"].clone(), r["more"].as_array().unwrap().len()), (json!(3), json!(false), 0));
|
||||
// nobody by that name, and bad input, are plain answers and errors
|
||||
let r = s.execute("faces.person", &json!({"name": "Nobody"})).unwrap();
|
||||
assert_eq!((r["total"].clone(), r["confirmed"].as_array().unwrap().len(), r["ready"].clone()), (json!(0), 0, json!(false)));
|
||||
for bad in [json!({}), json!({"name": ""}), json!({"name": " "}), json!({"name": 7})] {
|
||||
assert!(s.execute("faces.person", &bad).is_err(), "{bad}");
|
||||
}
|
||||
assert!(s.execute("faces.person", &json!({"name": "Ann", "more": u64::MAX})).is_ok());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
fn unit64(parts: &[(usize, f32)]) -> Vec<f32> {
|
||||
let mut v = vec![0.0f32; 64];
|
||||
for (i, x) in parts {
|
||||
v[*i] = *x;
|
||||
}
|
||||
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
v.iter().map(|x| x / n).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn many_faces_are_named_in_one_undo_step() {
|
||||
let d = temp("namemany");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
let (a, b) = two_photos(&s);
|
||||
let detected = |x: f64| Region { description: Some("Detected by YuNet 2023mar".into()), ..region(x, None) };
|
||||
set_regions(&mut s, a, vec![detected(0.1), region(0.5, None)]);
|
||||
set_regions(&mut s, b, vec![detected(0.3)]);
|
||||
let undo = s.undo.len();
|
||||
let r = s
|
||||
.execute(
|
||||
"faces.nameFaces",
|
||||
&json!({"faces": [{"photo": a.0, "index": 0}, {"photo": a.0, "index": 1}, {"photo": b.0, "index": 0}], "name": " Jane Doe "}),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!((r["named"].clone(), r["photos"].clone()), (json!(3), json!(2)));
|
||||
assert_eq!(s.undo.len(), undo + 1, "one step for the whole group");
|
||||
for (id, i) in [(a, 0), (a, 1), (b, 0)] {
|
||||
let reg = &s.catalog.photo(id).unwrap().meta.regions[i];
|
||||
assert_eq!(reg.name.as_deref(), Some("Jane Doe"));
|
||||
assert!(reg.description.is_none() || !reg.description.as_deref().unwrap().starts_with("Detected by"), "a named face is the user's now");
|
||||
}
|
||||
// undo restores every one of them together
|
||||
s.execute("edit.undo", &json!({})).unwrap();
|
||||
assert!(s.catalog.photo(a).unwrap().meta.regions.iter().chain(s.catalog.photo(b).unwrap().meta.regions.iter()).all(|r| r.name.is_none()));
|
||||
// an empty name clears; faces that do not exist are skipped, and with none left it is an error
|
||||
s.execute("faces.nameFaces", &json!({"faces": [{"photo": a.0, "index": 0}], "name": "Ann"})).unwrap();
|
||||
let r = s
|
||||
.execute(
|
||||
"faces.nameFaces",
|
||||
&json!({"faces": [{"photo": a.0, "index": 0}, {"photo": a.0, "index": 99}, {"photo": 987_654, "index": 0}], "name": null}),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(r["named"], 1);
|
||||
assert!(s.catalog.photo(a).unwrap().meta.regions[0].name.is_none());
|
||||
for bad in [
|
||||
json!({"name": "x"}),
|
||||
json!({"faces": "x", "name": "x"}),
|
||||
json!({"faces": [{"photo": a.0}], "name": "x"}),
|
||||
json!({"faces": [{"photo": 987_654, "index": 0}], "name": "x"}),
|
||||
json!({"faces": [{"photo": a.0, "index": 0}], "name": 7}),
|
||||
json!({"faces": [{"photo": a.0, "index": 0}], "name": "bad\u{7}name"}),
|
||||
json!({"faces": vec![json!({"photo": a.0, "index": 0}); 5001], "name": "x"}),
|
||||
] {
|
||||
assert!(s.execute("faces.nameFaces", &bad).is_err(), "{bad}");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unnamed_faces_come_with_look_alikes_together_and_the_name_the_named_ones_suggest() {
|
||||
let d = temp("unnamed");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
s.execute("faces.index", &json!({"budgetMs": 0})).unwrap();
|
||||
let (a, b) = two_photos(&s);
|
||||
let c = s.catalog.photos().map(|p| p.id).find(|id| *id != a && *id != b).unwrap();
|
||||
set_regions(&mut s, a, vec![region(0.1, Some("Ann")), region(0.5, None), region(0.7, None)]);
|
||||
set_regions(&mut s, b, vec![region(0.1, None), region(0.5, None)]);
|
||||
set_regions(&mut s, c, vec![region(0.3, None)]);
|
||||
// three faces like Ann, two like Bob (unnamed), one unembedded
|
||||
let put = |s: &mut Session, id: lightcraft_catalog::PhotoId, x: f64, v: Vec<f32>| s.faces.index.insert(id.0, ®ion(x, None).rect, v);
|
||||
put(&mut s, a, 0.1, unit64(&[(0, 1.0)])); // Ann, named
|
||||
put(&mut s, a, 0.5, unit64(&[(1, 1.0)])); // like Bob
|
||||
put(&mut s, b, 0.1, unit64(&[(0, 0.97), (2, 0.2)])); // like Ann
|
||||
put(&mut s, b, 0.5, unit64(&[(1, 0.98), (2, 0.2)])); // like Bob
|
||||
put(&mut s, c, 0.3, unit64(&[(0, 0.95), (3, 0.3)])); // like Ann
|
||||
// (a's last face, at 0.7, is not embedded yet)
|
||||
let r = s.execute("faces.unnamed", &json!({})).unwrap();
|
||||
assert_eq!((r["total"].clone(), r["ordered"].clone(), r["ready"].clone()), (json!(5), json!(true), json!(true)), "{r}");
|
||||
let faces = r["faces"].as_array().unwrap();
|
||||
let at = |f: &Value| (f["photo"].as_u64().unwrap(), f["index"].as_u64().unwrap());
|
||||
let order: Vec<(u64, u64)> = faces.iter().map(at).collect();
|
||||
// look-alikes are neighbours: the two like Bob are adjacent, and the three like Ann are adjacent; the unembedded one is last
|
||||
let pos = |x: (u64, u64)| order.iter().position(|o| *o == x).unwrap();
|
||||
assert_eq!(pos((a.0, 1)).abs_diff(pos((b.0, 1))), 1, "{order:?}");
|
||||
assert_eq!(order.last(), Some(&(a.0, 2)), "{order:?}");
|
||||
// a face the named ones recognise carries the name they suggest
|
||||
let like_ann = faces.iter().find(|f| at(f) == (b.0, 0)).unwrap();
|
||||
assert_eq!(like_ann["suggestion"]["name"], "Ann", "{like_ann}");
|
||||
assert!(faces.iter().find(|f| at(f) == (b.0, 1)).unwrap()["suggestion"].is_null(), "Bob is nobody the named faces know");
|
||||
// each face says which box to show it by; with no scan view it is its own box
|
||||
for k in ["x0", "y0", "x1", "y1"] {
|
||||
assert!((like_ann["view"][k].as_f64().unwrap() - like_ann["rect"][k].as_f64().unwrap()).abs() < 1e-3, "{like_ann}");
|
||||
}
|
||||
// a limit cuts the list, not the count
|
||||
let r = s.execute("faces.unnamed", &json!({"limit": 2})).unwrap();
|
||||
assert_eq!((r["total"].clone(), r["faces"].as_array().unwrap().len()), (json!(5), 2));
|
||||
// switched off, the faces are still listed, in photo order, without suggestions
|
||||
s.execute("faces.enable", &json!({"enabled": false})).unwrap();
|
||||
let r = s.execute("faces.unnamed", &json!({})).unwrap();
|
||||
assert_eq!((r["total"].clone(), r["ordered"].clone(), r["ready"].clone()), (json!(5), json!(false), json!(false)));
|
||||
assert!(r["faces"].as_array().unwrap().iter().all(|f| f["suggestion"].is_null()));
|
||||
for odd in [json!({"limit": 0}), json!({"limit": "x"}), json!({"limit": u64::MAX})] {
|
||||
assert!(s.execute("faces.unnamed", &odd).is_ok(), "{odd}");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_scanned_face_is_shown_by_the_same_kind_of_box() {
|
||||
let d = temp("view");
|
||||
let (mut s, _) = setup(&d, 64);
|
||||
s.execute("faces.index", &json!({"budgetMs": 0})).unwrap();
|
||||
let (a, _) = two_photos(&s);
|
||||
// a loosely drawn box (from some other tool) and the detector's tighter box of the same face
|
||||
let loose = lightcraft_geom::Rect { x0: 0.2, y0: 0.1, x1: 0.8, y1: 0.9 };
|
||||
let tight = lightcraft_geom::Rect { x0: 0.35, y0: 0.25, x1: 0.65, y1: 0.65 };
|
||||
let mut meta = s.catalog.photo(a).unwrap().meta.clone();
|
||||
meta.regions = vec![Region { rect: loose, kind: RegionKind::Face, name: Some("Ann".into()), description: None }];
|
||||
s.commit("setup", Op::SetMeta { id: a, meta: Box::new(meta) }).unwrap();
|
||||
// before the scan has looked at it, the face is shown by its own box
|
||||
assert_eq!(s.face_view(a, loose), loose);
|
||||
s.faces.index.insert_with_view(a.0, &loose, unit64(&[(0, 1.0)]), &tight);
|
||||
let shown = s.face_view(a, loose);
|
||||
assert!((shown.x0 - tight.x0).abs() < 1e-3 && (shown.y1 - tight.y1).abs() < 1e-3, "shown by the detector's box once known: {shown:?}");
|
||||
// a face nobody has looked at, or in another photo, is shown by its own box
|
||||
let other = lightcraft_geom::Rect { x0: 0.1, y0: 0.1, x1: 0.2, y1: 0.2 };
|
||||
assert_eq!(s.face_view(a, other), other);
|
||||
// the page of the person uses it too
|
||||
let r = s.execute("faces.person", &json!({"name": "Ann"})).unwrap();
|
||||
let view = &r["confirmed"][0]["view"];
|
||||
assert!((view["x0"].as_f64().unwrap() - tight.x0).abs() < 1e-3, "{r}");
|
||||
// and it survives a restart: the index file keeps it beside the embedding
|
||||
let (tag, path) = (s.faces.index.tag.clone(), d.join("view-cache.bin"));
|
||||
let mut written = crate::faces_index::Index::default();
|
||||
written.reset(&tag, 64, Some(path.clone()));
|
||||
written.insert_with_view(a.0, &loose, unit64(&[(0, 1.0)]), &tight);
|
||||
written.save().unwrap();
|
||||
drop(written);
|
||||
let mut again = crate::faces_index::Index::default();
|
||||
again.reset(&tag, 64, Some(path));
|
||||
let v = again.view(a.0, &loose).expect("the face is remembered");
|
||||
assert!((v.x0 - tight.x0).abs() < 1e-3 && (v.y0 - tight.y0).abs() < 1e-3, "{v:?}");
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
@@ -335,10 +335,15 @@ fn synthetic_dng(xmp: &str) -> Vec<u8> {
|
||||
|
||||
/// A tiny Bayer DNG written by our own DNG writer (optional embedded XMP, camera metadata).
|
||||
pub(crate) fn synthetic_dng_with(xmp: Option<&str>, metadata: lightcraft_meta::Metadata) -> Vec<u8> {
|
||||
use lightcraft_raw::*;
|
||||
let (w, h) = (32usize, 24usize);
|
||||
let cfa = Cfa::bayer("RGGB").unwrap();
|
||||
let data: Vec<u16> = (0..w * h).map(|i| 256 + ((i % w) * 300 + (i / w) * 200) as u16).collect();
|
||||
synthetic_dng_of(w, h, "RGGB", data, xmp, metadata)
|
||||
}
|
||||
|
||||
/// A synthetic Bayer DNG for denoise and cache tests.
|
||||
pub(crate) fn synthetic_dng_of(w: usize, h: usize, cfa: &str, data: Vec<u16>, xmp: Option<&str>, metadata: lightcraft_meta::Metadata) -> Vec<u8> {
|
||||
use lightcraft_raw::*;
|
||||
let cfa = Cfa::bayer(cfa).unwrap();
|
||||
let raw = RawImage {
|
||||
format: RawFormat::Dng,
|
||||
width: w,
|
||||
|
||||
Reference in New Issue
Block a user