use super::{
is_valid_pose, Effect, FragmentScale, Keyframe, Placement, Plane, ScanFields, SolveError,
SolveSession, PATCH_COUNT, PATCH_SIZE,
};
use crate::engine::attachments::{attachment_shift, resolve_target, Attachment};
use crate::engine::diagnostic::{Diagnostic, Severity, Value};
use crate::engine::keyframes::Retrieval;
use crate::engine::pose_graph::EdgeKind;
use crate::kernels::chrome::sticky_occlusions;
use crate::kernels::features::{match_features, Feature};
use crate::kernels::geometry::{js_round, Rect};
use crate::kernels::motion::{extract_patches, probe_scale, Gray, MatchPoints, Patch};
use crate::kernels::region::filter_features;
use crate::kernels::track::{
self, FragmentCauseGate, LoopVerdict, OdometryDecision, Rescue, ScreenFixed,
};
use std::collections::HashMap;
use std::sync::Arc;
const STATIC_CONFIDENCE: f64 = 0.3;
const BLIND_CONFIDENCE: f64 = 0.0;
const FRAGMENT_CONFIDENCE: f64 = 0.20;
const NONFINITE_CONFIDENCE: f64 = 0.0;
const LOOP_WEIGHT_CORRECTED: f64 = 6.0;
const LOOP_WEIGHT_RELINK: f64 = 5.0;
const LOOP_WEIGHT_CLOSURE: f64 = 4.0;
const CORRECTED_EDGE_WEIGHT: f64 = 0.05;
const DRIFT_DIFFERENCE: f64 = 0.12;
const REVISIT_AFTER: u32 = 3;
const TEXTURED_FEATURES: usize = 8;
const REGISTERED_FEATURES: usize = 3 * TEXTURED_FEATURES;
const PROBE_SCALES: [f64; 8] = [
1.1,
1.25,
1.5,
2.0,
1.0 / 1.1,
1.0 / 1.25,
1.0 / 1.5,
1.0 / 2.0,
];
pub(super) struct FrameCtx {
pub index: u32,
pub time: f64,
pub scan: ScanFields,
pub gray: Arc<Plane>,
pub features: Arc<Vec<Feature>>,
}
pub(super) enum Outcome {
Placed(Placement, bool),
Wait(Box<Step>, Vec<String>),
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Decision {
Tracked,
Static,
Blind,
Lost,
}
impl Decision {
fn tag(self) -> u8 {
match self {
Decision::Tracked => 0,
Decision::Static => 1,
Decision::Blind => 2,
Decision::Lost => 3,
}
}
}
enum Gate {
Start,
FixedInit,
Fixed,
Blind,
Lost,
Odometry,
}
enum Search {
Lost(Retrieval),
Revisit(Retrieval),
}
#[derive(Clone)]
struct Found {
keyframe: Arc<Keyframe>,
offset: (f64, f64),
confidence: f64,
ambiguous: bool,
support: u32,
unique: u32,
error: f64,
}
struct KeyLocals {
old_anchor: Option<Arc<Keyframe>>,
global: Option<Found>,
attached_from: Option<String>,
attach_match: Option<Found>,
corrected: Option<(f64, f64)>,
corrected_shift: (f64, f64),
edge_weight: f64,
}
pub(super) struct Step {
own: Vec<Feature>,
registration_rect: Rect,
full_overlap_rescue: bool,
textured: bool,
zoom_change: bool,
roi: Rect,
confidence: f64,
uncertain: bool,
relocalized: bool,
skip: bool,
ambiguous: bool,
lost: bool,
decision: Decision,
delta: (f64, f64),
via_anchor: Option<(f64, f64)>,
weak_step: bool,
step_error: f64,
rescue: Option<Rescue>,
keys: Option<KeyLocals>,
search: Option<Search>,
}
impl Step {
pub(super) fn waiting(&self) -> bool {
self.search.is_some()
}
}
fn js_max(a: f64, b: f64) -> f64 {
if a.is_nan() || b.is_nan() {
f64::NAN
} else if a == b {
if a == 0.0 && (a.is_sign_positive() || b.is_sign_positive()) {
0.0
} else {
a
}
} else {
a.max(b)
}
}
fn gray(p: &Plane) -> Gray<'_> {
Gray {
width: p.width,
height: p.height,
data: &p.data,
}
}
fn patches<'a>(
k: &'a Keyframe,
support: Option<Rect>,
exclude: Option<(&ScreenFixed, Rect, f64)>,
) -> Vec<Patch<'a>> {
let mut out = Vec::with_capacity(k.patches.len());
for p in &k.patches {
if support.is_some_and(|r| {
p.x < r.x
|| p.y < r.y
|| p.x + p.size as f64 > r.x + r.width
|| p.y + p.size as f64 > r.y + r.height
}) {
continue;
}
if exclude.is_some_and(|(cells, rect, f)| {
let half = p.size as f64 / 2.0;
cells.excludes((rect.x + p.x + half) / f, (rect.y + p.y + half) / f)
}) {
continue;
}
out.push(Patch {
x: p.x as i32,
y: p.y as i32,
size: p.size,
data: &p.data,
});
}
out
}
fn pad(n: u32) -> String {
format!("{n:0>10}")
}
impl SolveSession {
pub(super) fn step_region(
&mut self,
i: usize,
ctx: &FrameCtx,
labels: &[u8],
) -> Result<Outcome, SolveError> {
let mut s = self.setup(i, ctx);
let registration = self.registration_plane(labels);
let registration_labels = registration.as_deref().map_or(labels, Vec::as_slice);
match self.gate(i, ctx, &s) {
Gate::Odometry => self.odometry_step(i, ctx, &mut s, registration_labels),
Gate::Start => {
self.states[i].started = true;
self.states[i].canvas_id = String::new();
self.new_canvas(i, ctx.time);
}
Gate::FixedInit => self.new_canvas(i, ctx.time),
Gate::Fixed => {}
Gate::Blind => s.decision = Decision::Blind,
Gate::Lost => s.decision = Decision::Lost,
}
self.reacquire_step(i, &mut s);
if self.regions[i].moving {
match s.decision {
Decision::Tracked => self.apply_tracked(i, ctx, &mut s),
Decision::Static => self.apply_static(i, ctx, &mut s),
Decision::Blind => self.apply_blind(i, ctx, &mut s),
Decision::Lost => {
s.lost = true;
let exclude = self.states[i].anchor.as_ref().map(|a| a.id.clone());
let retrieval = self.retrieve(i, ctx.index, &s.own, exclude.as_deref());
if !retrieval.candidates.is_empty() {
let ids = retrieval.candidates.clone();
s.search = Some(Search::Lost(retrieval));
return Ok(Outcome::Wait(Box::new(s), ids));
}
self.apply_lost(i, ctx, &mut s, None);
}
}
}
self.after_decision(i, ctx, s)
}
pub(super) fn resume_region(
&mut self,
i: usize,
ctx: &FrameCtx,
mut s: Step,
rows: Vec<Option<Keyframe>>,
) -> Result<Outcome, SolveError> {
match s.search.take() {
Some(Search::Lost(retrieval)) => {
let found = self.evaluate(i, ctx, &s, &retrieval, rows);
self.apply_lost(i, ctx, &mut s, found);
self.after_decision(i, ctx, s)
}
Some(Search::Revisit(retrieval)) => {
let found = self.evaluate(i, ctx, &s, &retrieval, rows);
self.keyframe_rest(i, ctx, s, found)
}
None => Err(SolveError::BadInput("no keyframes were asked for")),
}
}
fn patch_exclusion(&self, i: usize) -> Option<(&ScreenFixed, Rect, f64)> {
self.states[i]
.screen_fixed
.as_ref()
.filter(|cells| cells.any())
.map(|cells| (cells, self.regions[i].shape.rect, self.factor as f64))
}
fn registration_plane(&mut self, labels: &[u8]) -> Option<Arc<Vec<u8>>> {
if !self
.states
.iter()
.any(|s| s.screen_fixed.as_ref().is_some_and(ScreenFixed::any))
{
self.registration_labels = None;
return None;
}
let current = !self.registration_stale
&& self
.registration_labels
.as_ref()
.is_some_and(|(_, base)| base.as_slice() == labels);
if !current {
let (w, h, f) = (self.width, self.height, self.factor as f64);
let mut plane = labels.to_vec();
for (r, state) in self.regions.iter().zip(&self.states) {
let Some(cells) = state.screen_fixed.as_ref().filter(|c| c.any()) else {
continue;
};
let rect = r.shape.rect;
let (x0, y0) = (rect.x.max(0.0) as usize, rect.y.max(0.0) as usize);
let x1 = ((rect.x + rect.width).ceil().max(0.0) as usize).min(w);
let y1 = ((rect.y + rect.height).ceil().max(0.0) as usize).min(h);
for y in y0..y1 {
let ay = y as f64 / f;
for x in x0..x1 {
let k = y * w + x;
if plane[k] == r.code && cells.excludes(x as f64 / f, ay) {
plane[k] = super::UNREGISTERED;
}
}
}
}
self.registration_labels = Some((Arc::new(plane), labels.to_vec()));
self.registration_stale = false;
}
self.registration_labels
.as_ref()
.map(|(plane, _)| plane.clone())
}
fn patch_support(&self, i: usize, s: &Step) -> Option<Rect> {
let full = self.regions[i].shape.rect;
(s.registration_rect != full).then_some(Rect {
x: s.registration_rect.x - full.x,
y: s.registration_rect.y - full.y,
width: s.registration_rect.width,
height: s.registration_rect.height,
})
}
fn setup(&self, i: usize, ctx: &FrameCtx) -> Step {
let r = &self.regions[i];
let state = &self.states[i];
let f = self.factor as f64;
let rect = r.shape.rect;
let roi = Rect {
x: rect.x / f,
y: rect.y / f,
width: rect.width / f,
height: rect.height / f,
};
let own = filter_features(
&r.shape,
&ctx.features,
f,
self.width as f64,
self.height as f64,
);
let (registration_rect, own) =
track::registration_features(rect, self.states[i].registration_support, own, f);
let own = match state.screen_fixed.as_ref().filter(|c| c.any()) {
Some(cells) => {
let kept: Vec<Feature> = own
.iter()
.filter(|p| !cells.excludes(p.x as f64, p.y as f64))
.cloned()
.collect();
if kept.len() >= REGISTERED_FEATURES {
kept
} else {
own
}
}
None => own,
};
let textured = own.len() >= TEXTURED_FEATURES;
let prior: Vec<MatchPoints> = if r.moving {
let previous = state.previous_features.as_deref().unwrap_or(&[]);
match_features(previous, &own, true)
.iter()
.map(|m| {
let (a, b) = (&previous[m.a as usize], &own[m.b as usize]);
MatchPoints {
ax: a.x as f64,
ay: a.y as f64,
bx: b.x as f64,
by: b.y as f64,
unique: m.unique,
}
})
.collect()
} else {
Vec::new()
};
let zoom = track::region_zoom(r.moving, &prior);
let zoom_change = track::zoom_changed(zoom, ctx.scan.zoom);
Step {
own,
registration_rect,
full_overlap_rescue: false,
textured,
zoom_change,
roi,
confidence: if r.unassigned { 0.2 } else { 1.0 },
uncertain: r.unassigned,
relocalized: false,
skip: false,
ambiguous: false,
lost: false,
decision: Decision::Tracked,
delta: (0.0, 0.0),
via_anchor: None,
weak_step: false,
step_error: f64::INFINITY,
rescue: None,
keys: None,
search: None,
}
}
fn gate(&self, i: usize, ctx: &FrameCtx, s: &Step) -> Gate {
let (moving, state) = (self.regions[i].moving, &self.states[i]);
let has_previous = self.previous_slot.is_some() && self.previous_gray.is_some();
if moving && !state.started {
if s.textured {
Gate::Start
} else {
Gate::Blind
}
} else if !moving {
if ctx.index == 0 {
Gate::FixedInit
} else {
Gate::Fixed
}
} else if !s.textured {
Gate::Blind
} else if state.blind || s.zoom_change || !has_previous {
Gate::Lost
} else {
Gate::Odometry
}
}
fn odometry_step(&mut self, i: usize, ctx: &FrameCtx, s: &mut Step, labels: &[u8]) {
let r = &self.regions[i];
let state = &self.states[i];
let previous_gray = self
.previous_gray
.as_ref()
.expect("gated on a previous frame");
let previous_slot = self.previous_slot.expect("gated on a previous frame");
let inputs = track::OdometryInputs {
f: self.factor as f64,
radius: self.radius,
mask: Some((labels, r.code)),
roi: Rect {
x: s.registration_rect.x / self.factor as f64,
y: s.registration_rect.y / self.factor as f64,
width: s.registration_rect.width / self.factor as f64,
height: s.registration_rect.height / self.factor as f64,
},
rect: s.registration_rect,
region: Some(&r.shape),
image_width: self.width as f64,
image_height: self.height as f64,
previous: &self.frames[previous_slot],
current: &self.frames[self.current_slot],
previous_gray: gray(previous_gray),
g: gray(&ctx.gray),
velocity: state.velocity,
previous_features: state.previous_features.as_deref().unwrap_or(&[]),
own_features: &s.own,
confidence: s.confidence,
};
#[allow(unused_mut)]
let mut est = track::odometry(inputs);
if s.registration_rect != r.shape.rect
&& est.decision == OdometryDecision::Lost
&& est.rescue.is_none()
{
let own = filter_features(
&r.shape,
&ctx.features,
self.factor as f64,
self.width as f64,
self.height as f64,
);
let previous = filter_features(
&r.shape,
self.previous_features
.as_deref()
.map(Vec::as_slice)
.unwrap_or(&[]),
self.factor as f64,
self.width as f64,
self.height as f64,
);
let full = track::odometry(track::OdometryInputs {
roi: s.roi,
rect: r.shape.rect,
previous_features: &previous,
own_features: &own,
..inputs
});
if full.decision == OdometryDecision::Tracked && full.delta != (0.0, 0.0) {
s.full_overlap_rescue = true;
est.rescue = Some(Rescue {
phase_support: false,
delta: full.delta,
confidence: full.confidence,
ambiguous: full.ambiguous,
weak_step: full.weak_step,
step_error: full.step_error,
});
} else if full.rescue.as_ref().is_some_and(|r| r.delta != (0.0, 0.0)) {
s.full_overlap_rescue = true;
est.rescue = full.rescue;
}
}
#[cfg(test)]
if let Some(seam) = self.seams.odometry.as_mut() {
self.seams.reacquire_fails |= seam(ctx.index, &mut est);
}
s.decision = match est.decision {
OdometryDecision::Tracked => Decision::Tracked,
OdometryDecision::Static => Decision::Static,
OdometryDecision::Lost => Decision::Lost,
};
s.delta = est.delta;
s.ambiguous = est.ambiguous;
s.weak_step = est.weak_step;
s.confidence = est.confidence;
s.step_error = est.step_error;
s.rescue = est.rescue;
if let Some(change) = est.content_change {
let canvas_id = self.states[i].canvas_id.clone();
self.diagnostic(
Diagnostic::new("PARTIAL_CONTENT_CHANGE", Severity::Info)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.param("percent", js_round((1.0 - change.agreement) * 100.0) as f64),
);
}
}
fn reacquire_step(&mut self, i: usize, s: &mut Step) {
let Some(anchor) = self.states[i].anchor.clone() else {
return;
};
if !(s.decision == Decision::Lost && !s.zoom_change && self.regions[i].moving) {
return;
}
let models = track::reacquire_hypotheses(&anchor.features, &s.own);
if models.is_empty() {
return;
}
let (rect, f, radius, w, h) = (
self.regions[i].shape.rect,
self.factor as f64,
self.radius,
self.width,
self.height,
);
let anchor_patches = patches(&anchor, self.patch_support(i, s), self.patch_exclusion(i));
let native = Gray {
width: w,
height: h,
data: self.native(),
};
let top = track::reacquire_refine(&models, &anchor_patches, native, rect, f, radius);
#[cfg(test)]
#[allow(unused_mut)]
let mut top = top.filter(|_| !std::mem::take(&mut self.seams.reacquire_fails));
#[cfg(test)]
if let Some(seam) = self.seams.reacquire.as_mut() {
seam(&mut top);
}
if let Some(top) = top {
let pose = self.states[i].pose;
let implied = (
anchor.x + top.x as f64 - pose.0,
anchor.y + top.y as f64 - pose.1,
);
if s.rescue
.as_ref()
.is_some_and(|r| track::reacquisition_contradicts_rescue(implied, r.delta))
{
return;
}
s.decision = Decision::Tracked;
s.via_anchor = Some((top.x as f64, top.y as f64));
s.ambiguous = top.ambiguous;
s.confidence = top.confidence;
}
}
fn apply_tracked(&mut self, i: usize, ctx: &FrameCtx, s: &mut Step) {
let rect = self.regions[i].shape.rect;
let support = self.patch_support(i, s);
if let Some((vx, vy)) = s.via_anchor {
let state = &mut self.states[i];
let anchor = state
.anchor
.as_ref()
.expect("re-acquired through the anchor");
state.pose = (anchor.x + vx, anchor.y + vy);
state.velocity = (0.0, 0.0);
} else if ctx.index > 0 && self.previous_slot.is_some() {
let anchor = self.states[i].anchor.clone();
let anchor_patches = anchor
.as_ref()
.map(|a| patches(a, support, self.patch_exclusion(i)));
let state = &mut self.states[i];
state.pose = (state.pose.0 + s.delta.0, state.pose.1 + s.delta.1);
state.velocity = s.delta;
if let (Some(anchor), Some(anchor_patches)) = (anchor.as_ref(), anchor_patches) {
if ctx.scan.difference >= DRIFT_DIFFERENCE {
let (pose, radius, w, h) = (state.pose, self.radius, self.width, self.height);
let native = Gray {
width: w,
height: h,
data: self.native(),
};
#[allow(unused_mut)]
let mut drift = track::drift_correction(
&anchor_patches,
native,
rect,
(anchor.x, anchor.y),
pose,
radius,
);
#[cfg(test)]
if let Some(seam) = self.seams.drift.as_mut() {
seam(&mut drift);
}
if let Some(d) = drift {
self.states[i].pose = (d.x, d.y);
s.confidence = js_max(s.confidence, d.confidence_floor);
}
}
}
}
s.uncertain = track::uncertainty(s.confidence, s.ambiguous, s.weak_step);
let (canvas_id, pose) = (self.states[i].canvas_id.clone(), self.states[i].pose);
if s.weak_step {
self.diagnostic(
Diagnostic::new("THIN_OVERLAP_STEP", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(s.confidence),
);
}
if s.uncertain {
let mut diagnostic = Diagnostic::new("LOW_CONFIDENCE_PLACEMENT", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(s.confidence)
.region(Rect {
x: pose.0 + rect.x,
y: pose.1 + rect.y,
width: rect.width,
height: rect.height,
})
.param("ambiguous", s.ambiguous as u8 as f64);
if s.registration_rect != rect {
let mut detail = vec![
("registrationScreenRect", Value::rect(s.registration_rect)),
("outputScreenRect", Value::rect(rect)),
];
if s.full_overlap_rescue && s.via_anchor.is_none() && !s.relocalized {
detail.push(("fullRegionFallback", Value::Bool(true)));
}
diagnostic = diagnostic.detail(Value::Object(detail));
}
self.diagnostic(diagnostic);
}
}
fn apply_static(&mut self, i: usize, ctx: &FrameCtx, s: &mut Step) {
self.states[i].velocity = (0.0, 0.0);
s.confidence = STATIC_CONFIDENCE;
s.uncertain = true;
let canvas_id = self.states[i].canvas_id.clone();
self.diagnostic(
Diagnostic::new("LOW_CONFIDENCE_PLACEMENT", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(s.confidence),
);
}
fn apply_blind(&mut self, i: usize, ctx: &FrameCtx, s: &mut Step) {
s.skip = true;
s.confidence = BLIND_CONFIDENCE;
s.uncertain = true;
self.states[i].velocity = (0.0, 0.0);
let canvas_id = self.states[i].canvas_id.clone();
self.diagnostic(
Diagnostic::new("UNOBSERVABLE_FRAME", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id),
);
}
fn apply_lost(&mut self, i: usize, ctx: &FrameCtx, s: &mut Step, found: Option<Found>) {
let verdict = found.as_ref().is_some_and(|m| {
track::relocalize_verdict(
Some(track::MatchInfo {
ambiguous: m.ambiguous,
confidence: m.confidence,
}),
s.zoom_change,
) && !self.relocalization_contradicts_rescue(i, s, m)
});
if let (true, Some(m)) = (verdict, found) {
let canvas = resolve_target(&self.attachments, &m.keyframe.canvas_id).to_string();
let shift = attachment_shift(&self.attachments, &m.keyframe.canvas_id);
let state = &mut self.states[i];
state.canvas_id = canvas;
state.pose = track::target_pose((m.keyframe.x, m.keyframe.y), m.offset, shift);
state.last_node = None;
state.anchor = Some(m.keyframe.clone());
state.velocity = (0.0, 0.0);
s.confidence = m.confidence;
s.relocalized = true;
let canvas_id = state.canvas_id.clone();
self.diagnostic(
Diagnostic::new("RELOCALIZED", Severity::Info)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(s.confidence)
.detail(Value::Object(vec![
("anchorFrame", Value::Number(m.keyframe.frame as f64)),
("support", Value::Number(m.support as f64)),
("unique", Value::Number(m.unique as f64)),
("error", Value::Number(m.error)),
("offset", Value::point(m.offset.0, m.offset.1)),
])),
);
} else if let (Some(rescue), false) = (s.rescue.as_ref(), s.zoom_change) {
s.lost = false;
s.decision = Decision::Tracked;
s.delta = rescue.delta;
s.confidence = rescue.confidence;
s.ambiguous = rescue.ambiguous;
s.weak_step = rescue.weak_step;
s.step_error = rescue.step_error;
if rescue.phase_support && !self.states[i].phase_support {
self.states[i].registration_support = track::registration_support(
self.regions[i].shape.rect,
&s.own,
self.factor as f64,
);
self.states[i].phase_support = true;
let (rect, own) = track::registration_features(
self.regions[i].shape.rect,
self.states[i].registration_support,
std::mem::take(&mut s.own),
self.factor as f64,
);
s.registration_rect = rect;
s.own = own;
}
self.apply_tracked(i, ctx, s);
} else {
let scale: FragmentScale = match track::fragment_cause_gate(
s.zoom_change,
self.previous_gray.is_some(),
self.states[i].blind,
) {
FragmentCauseGate::ZoomChange => Some((ctx.scan.zoom, 0.0)),
FragmentCauseGate::ProbeScale => probe_scale(
gray(
self.previous_gray
.as_ref()
.expect("gated on a previous gray"),
),
gray(&ctx.gray),
&s.own,
Some(s.roi),
&PROBE_SCALES,
),
FragmentCauseGate::None => None,
};
let state = &mut self.states[i];
state.fragment += 1;
state.pose = (0.0, 0.0);
state.last_node = None;
state.anchor = None;
state.velocity = (0.0, 0.0);
self.new_canvas(i, ctx.time);
s.confidence = FRAGMENT_CONFIDENCE;
s.uncertain = true;
let canvas_id = self.states[i].canvas_id.clone();
let d = match scale {
Some((scale, error)) => Diagnostic::new("SCALE_CHANGE_FRAGMENT", Severity::Warning)
.param("scale", scale)
.detail(Value::Object(vec![
("scale", Value::Number(scale)),
("error", Value::Number(error)),
])),
None => Diagnostic::new("UNPLACED_FRAGMENT", Severity::Warning),
};
self.diagnostic(
d.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(s.confidence),
);
}
}
fn relocalization_contradicts_rescue(&self, i: usize, s: &Step, m: &Found) -> bool {
let Some(rescue) = s.rescue.as_ref() else {
return false;
};
let state = &self.states[i];
if resolve_target(&self.attachments, &m.keyframe.canvas_id) != state.canvas_id {
return false;
}
let shift = attachment_shift(&self.attachments, &m.keyframe.canvas_id);
let target = track::target_pose((m.keyframe.x, m.keyframe.y), m.offset, shift);
track::reacquisition_contradicts_rescue(
(target.0 - state.pose.0, target.1 - state.pose.1),
rescue.delta,
)
}
fn recover_nonfinite_pose(&mut self, i: usize, ctx: &FrameCtx, s: &mut Step) {
let (x, y) = self.states[i].pose;
if is_valid_pose(x, y) {
return;
}
self.diagnostic(Diagnostic::new("NONFINITE_POSE", Severity::Error).at(ctx.index, ctx.time));
let state = &mut self.states[i];
state.fragment += 1;
state.pose = (0.0, 0.0);
state.last_node = None;
state.anchor = None;
s.confidence = NONFINITE_CONFIDENCE;
s.uncertain = true;
self.new_canvas(i, ctx.time);
}
fn after_decision(
&mut self,
i: usize,
ctx: &FrameCtx,
mut s: Step,
) -> Result<Outcome, SolveError> {
self.recover_nonfinite_pose(i, ctx, &mut s);
if s.skip {
return Ok(self.close_region(i, ctx, s));
}
let (r, state) = (&self.regions[i], &self.states[i]);
let needs = track::needs_keyframe(
r.moving,
state.anchor.as_ref().map(|a| (a.x, a.y)),
state.pose,
state.last_node.as_ref().map(|n| n.1 as i32),
(r.shape.rect.width, r.shape.rect.height),
ctx.index as i32,
ctx.scan.difference,
);
if !needs {
return Ok(self.close_region(i, ctx, s));
}
let old_anchor = state.anchor.clone();
s.keys = Some(KeyLocals {
old_anchor: old_anchor.clone(),
global: None,
attached_from: None,
attach_match: None,
corrected: None,
corrected_shift: (0.0, 0.0),
edge_weight: track::odometry_weight(s.weak_step),
});
if !s.lost && !s.skip && r.moving && ctx.index > REVISIT_AFTER {
let exclude = old_anchor.as_ref().map(|a| a.id.clone());
let retrieval = self.retrieve(i, ctx.index, &s.own, exclude.as_deref());
if !retrieval.candidates.is_empty() {
let ids = retrieval.candidates.clone();
s.search = Some(Search::Revisit(retrieval));
return Ok(Outcome::Wait(Box::new(s), ids));
}
}
self.keyframe_rest(i, ctx, s, None)
}
fn evaluate(
&mut self,
i: usize,
ctx: &FrameCtx,
s: &Step,
retrieval: &Retrieval,
rows: Vec<Option<Keyframe>>,
) -> Option<Found> {
let keyframes: Vec<Arc<Keyframe>> = rows
.into_iter()
.flatten()
.filter(|k| retrieval.admits(k.frame))
.map(Arc::new)
.collect();
if keyframes.is_empty() {
return None;
}
let mut canonical: HashMap<&str, (String, (f64, f64))> = HashMap::new();
let mut interned: HashMap<String, u32> = HashMap::new();
let mut resolved = Vec::with_capacity(keyframes.len());
for k in &keyframes {
let (canvas, shift) = canonical
.entry(k.canvas_id.as_str())
.or_insert_with(|| {
(
resolve_target(&self.attachments, &k.canvas_id).to_string(),
attachment_shift(&self.attachments, &k.canvas_id),
)
})
.clone();
let next = interned.len() as u32;
let idx = *interned.entry(canvas).or_insert(next);
resolved.push((idx, shift));
}
let patch_lists: Vec<Vec<Patch>> = keyframes
.iter()
.map(|k| patches(k, self.patch_support(i, s), self.patch_exclusion(i)))
.collect();
let candidates: Vec<track::CandidateKeyframe> = keyframes
.iter()
.zip(&patch_lists)
.zip(&resolved)
.map(
|((k, p), &(canonical_idx, (dx, dy)))| track::CandidateKeyframe {
features: &k.features,
gray: gray(&k.gray),
patches: p,
x: k.x,
y: k.y,
canonical_idx,
dx,
dy,
},
)
.collect();
let audited = track::evaluate_candidates_audit(
&candidates,
&s.own,
gray(&ctx.gray),
s.roi,
self.factor as f64,
);
if audited.is_empty() {
return None;
}
let (rect, radius, w, h) = (
self.regions[i].shape.rect,
self.radius,
self.width,
self.height,
);
let native = Gray {
width: w,
height: h,
data: self.native(),
};
let refined =
track::evaluate_candidates_refine(&audited, &candidates, native, rect, radius);
let chosen = track::select_candidate(&refined, &candidates)?;
Some(Found {
keyframe: keyframes[chosen.keyframe_index].clone(),
offset: (chosen.x as f64, chosen.y as f64),
confidence: chosen.confidence,
ambiguous: chosen.ambiguous,
support: chosen.support,
unique: chosen.unique,
error: chosen.error,
})
}
fn keyframe_rest(
&mut self,
i: usize,
ctx: &FrameCtx,
mut s: Step,
found: Option<Found>,
) -> Result<Outcome, SolveError> {
let mut keys = s.keys.take().expect("the keyframe half started");
keys.global = found.filter(|m| {
is_valid_pose(m.keyframe.x, m.keyframe.y) && is_valid_pose(m.offset.0, m.offset.1)
});
self.apply_attachment(i, ctx, &mut keys);
self.apply_thin_overlap_correction(i, ctx, &s, &mut keys);
let (canvas_id, pose, previous) = {
let state = &self.states[i];
(
state.canvas_id.clone(),
state.pose,
state.last_node.as_ref().map(|n| n.0.clone()),
)
};
let (node, writes) = self.graph.add(
&canvas_id,
ctx.index,
pose,
previous.as_deref(),
keys.edge_weight,
)?;
self.graph_writes(writes);
if let (Some(_), Some(global)) = (keys.corrected, &keys.global) {
let writes = self.graph.connect(
&global.keyframe.node,
&node.id,
global.offset.0 + keys.corrected_shift.0,
global.offset.1 + keys.corrected_shift.1,
LOOP_WEIGHT_CORRECTED,
EdgeKind::Loop,
)?;
self.graph_writes(writes);
}
self.connect_relink_or_loop_closure(i, ctx, &s, &keys, &node.id)?;
let last = self.graph.get(&node.id).map(|n| (n.id.clone(), n.frame));
self.states[i].last_node = last;
self.mint_keyframe(i, ctx, &s, &node.id);
s.keys = Some(keys);
Ok(self.close_region(i, ctx, s))
}
fn apply_attachment(&mut self, i: usize, ctx: &FrameCtx, keys: &mut KeyLocals) {
let Some(global) = keys.global.clone() else {
return;
};
let target = resolve_target(&self.attachments, &global.keyframe.canvas_id).to_string();
let shift = attachment_shift(&self.attachments, &global.keyframe.canvas_id);
let canvas_id = self.states[i].canvas_id.clone();
#[allow(unused_mut)]
let mut verdict = track::attach_verdict(
Some(track::GlobalMatch {
keyframe: (global.keyframe.x, global.keyframe.y),
offset: global.offset,
ambiguous: global.ambiguous,
confidence: global.confidence,
}),
target == canvas_id,
shift,
);
#[cfg(test)]
if let Some(seam) = self.seams.attach.as_mut() {
seam(&mut verdict);
}
let Some(pose) = verdict.filter(|p| is_valid_pose(p.0, p.1)) else {
return;
};
let state_pose = self.states[i].pose;
let attachment = Attachment {
id: canvas_id.clone(),
target: target.clone(),
dx: pose.0 - state_pose.0,
dy: pose.1 - state_pose.1,
frame: ctx.index,
confidence: global.confidence,
node: global.keyframe.node.clone(),
};
self.attachments
.insert(canvas_id.clone(), attachment.clone());
self.effects.push(Effect::Attachment(attachment.clone()));
if let Some(meta) = self.canvases.get_mut(&canvas_id) {
meta.attached_to = Some(target.clone());
let meta = meta.clone();
self.effects.push(Effect::Canvas {
meta,
created: false,
});
}
self.diagnostic(
Diagnostic::new("FRAGMENT_ATTACHED", Severity::Info)
.at(ctx.index, ctx.time)
.canvas(&attachment.target)
.confidence(global.confidence)
.detail(Value::Object(vec![
("id", Value::Text(attachment.id.clone())),
("target", Value::Text(attachment.target.clone())),
("dx", Value::Number(attachment.dx)),
("dy", Value::Number(attachment.dy)),
("frame", Value::Number(attachment.frame as f64)),
("confidence", Value::Number(attachment.confidence)),
("node", Value::Text(attachment.node.clone())),
])),
);
keys.attached_from = Some(canvas_id);
keys.attach_match = Some(global);
let state = &mut self.states[i];
state.canvas_id = target;
state.pose = pose;
state.last_node = None;
state.velocity = (0.0, 0.0);
}
fn apply_thin_overlap_correction(
&mut self,
i: usize,
ctx: &FrameCtx,
s: &Step,
keys: &mut KeyLocals,
) {
let Some(global) = keys.global.clone() else {
return;
};
let state = &self.states[i];
#[allow(unused_mut)]
let mut eligible = track::thin_overlap_eligible(
s.weak_step,
state.weak,
global.ambiguous,
global.confidence,
global.error,
);
#[cfg(test)]
{
eligible |= self.seams.thin.is_some();
}
if keys.attached_from.is_some()
|| resolve_target(&self.attachments, &global.keyframe.canvas_id) != state.canvas_id
|| !eligible
{
return;
}
let shift = attachment_shift(&self.attachments, &global.keyframe.canvas_id);
let canonical = (global.keyframe.x + shift.0, global.keyframe.y + shift.1);
#[allow(unused_mut)]
let mut correction = track::thin_overlap_correction(canonical, global.offset, state.pose);
#[cfg(test)]
if let Some(seam) = self.seams.thin.as_mut() {
seam(&mut correction);
}
let Some((target, discrepancy)) = correction.filter(|(t, _)| is_valid_pose(t.0, t.1))
else {
return;
};
let state = &self.states[i];
keys.corrected = Some(target);
keys.corrected_shift = shift;
keys.edge_weight = CORRECTED_EDGE_WEIGHT;
let (canvas_id, from) = (state.canvas_id.clone(), state.pose);
self.diagnostic(
Diagnostic::new("TRAJECTORY_CORRECTED", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(global.confidence)
.param("discrepancy", discrepancy)
.detail(Value::Object(vec![
("discrepancy", Value::Number(discrepancy)),
("from", Value::point(from.0, from.1)),
("to", Value::point(target.0, target.1)),
("revisitError", Value::Number(global.error)),
("stepError", Value::Number(s.step_error)),
("weakStep", Value::Bool(s.weak_step)),
])),
);
let state = &mut self.states[i];
state.pose = target;
state.weak = false;
}
fn connect_relink_or_loop_closure(
&mut self,
i: usize,
ctx: &FrameCtx,
s: &Step,
keys: &KeyLocals,
node: &str,
) -> Result<(), SolveError> {
let (canvas_id, pose) = (self.states[i].canvas_id.clone(), self.states[i].pose);
if (s.relocalized && keys.old_anchor.is_some()) || keys.attached_from.is_some() {
let writes = self.graph.unpin(node)?;
self.graph_writes(writes);
let link = match &keys.attach_match {
Some(m) => m.keyframe.clone(),
None => keys.old_anchor.clone().expect("relocalized with an anchor"),
};
let writes = self.graph.connect(
&link.node,
node,
pose.0 - link.x,
pose.1 - link.y,
LOOP_WEIGHT_RELINK,
EdgeKind::Loop,
)?;
self.graph_writes(writes);
} else if let Some(global) = &keys.global {
if resolve_target(&self.attachments, &global.keyframe.canvas_id) != canvas_id {
return Ok(());
}
let shift = attachment_shift(&self.attachments, &global.keyframe.canvas_id);
let (verdict, discrepancy) = track::loop_closure_verdict(
(global.keyframe.x, global.keyframe.y),
global.offset,
global.ambiguous,
global.confidence,
shift,
pose,
);
match verdict {
LoopVerdict::Closure => {
let writes = self.graph.connect(
&global.keyframe.node,
node,
global.offset.0 + shift.0,
global.offset.1 + shift.1,
LOOP_WEIGHT_CLOSURE,
EdgeKind::Loop,
)?;
self.graph_writes(writes);
self.diagnostic(
Diagnostic::new("LOOP_CLOSURE", Severity::Info)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(global.confidence),
);
}
LoopVerdict::Inconsistent => self.diagnostic(
Diagnostic::new("INCONSISTENT_LOOP_REJECTED", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.confidence(global.confidence)
.param("discrepancy", discrepancy),
),
LoopVerdict::Ambiguous => self.diagnostic(
Diagnostic::new("AMBIGUOUS_LOOP", Severity::Warning)
.at(ctx.index, ctx.time)
.canvas(&canvas_id),
),
LoopVerdict::None => {}
}
}
Ok(())
}
fn mint_keyframe(&mut self, i: usize, ctx: &FrameCtx, s: &Step, node: &str) {
let f = self.factor as f64;
let (moving, rect, id) = (
self.regions[i].moving,
self.regions[i].shape.rect,
self.regions[i].id.clone(),
);
let patches = if moving {
let points: Vec<(f64, f64)> = s
.own
.iter()
.map(|p| (p.x as f64 - s.roi.x, p.y as f64 - s.roi.y))
.collect();
let (w, h) = (self.width, self.height);
let native = self.native();
extract_patches(native, w, h, rect, &points, f, PATCH_COUNT, PATCH_SIZE)
.into_iter()
.map(|p| super::KeyPatch {
x: p.x,
y: p.y,
size: p.size,
data: p.data,
})
.collect()
} else {
Vec::new()
};
let state = &self.states[i];
let keyframe = Arc::new(Keyframe {
id: format!("{id}/{}", pad(ctx.index)),
node: node.to_string(),
canvas_id: state.canvas_id.clone(),
layer: id.clone(),
frame: ctx.index,
features: s.own.clone(),
gray: ctx.gray.clone(),
x: state.pose.0,
y: state.pose.1,
scale_x: f,
scale_y: f,
patches,
});
self.states[i].anchor = Some(keyframe.clone());
if moving {
let descriptors: Vec<_> = keyframe.features.iter().map(|f| f.descriptor).collect();
self.effects.push(Effect::Keyframe(keyframe.clone()));
self.index
.add(&keyframe.id, &id, keyframe.frame, &descriptors);
}
}
fn close_region(&mut self, i: usize, ctx: &FrameCtx, mut s: Step) -> Outcome {
self.recover_nonfinite_pose(i, ctx, &mut s);
let (moving, layer, region) = (
self.regions[i].moving,
self.regions[i].id.clone(),
self.regions[i]
.shape
.crop
.unwrap_or(self.regions[i].shape.rect),
);
let mut observed = false;
if moving && !s.skip {
let slot = self.voting_slot[i].expect("a moving region has a voting slot");
let canvas_id = self.states[i].canvas_id.clone();
let canvas = self.canvas_slot(&canvas_id);
let pose = self.states[i].pose;
self.voting
.observe(slot, canvas, pose.0, pose.1, &ctx.gray.data);
observed = true;
}
if moving
&& !s.skip
&& s.decision == Decision::Tracked
&& s.via_anchor.is_none()
&& !s.relocalized
{
if let Some(previous) = self.previous_gray.clone() {
let (f, factor) = (self.factor as f64, self.factor as i64);
let (w, h) = (self.width as f64, self.height as f64);
let shape = &self.regions[i].shape;
if let Some(cells) = self.states[i].screen_fixed.as_mut() {
let inside = |x: i64, y: i64| {
shape.contains(
(x * factor + factor / 2) as f64,
(y * factor + factor / 2) as f64,
w,
h,
)
};
if cells.observe(
&previous.data,
&ctx.gray.data,
previous.width,
previous.height,
(s.delta.0 / f, s.delta.1 / f),
inside,
) {
self.registration_stale = true;
}
}
}
}
let has_previous = self.previous_slot.is_some();
let canvas_id = self.states[i].canvas_id.clone();
let occlusions = match self.previous_slot {
Some(previous) if track::occlusion_eligible(has_previous, moving, s.decision.tag()) => {
let carry: Vec<Rect> = self
.previous_plan
.as_ref()
.and_then(|plan| {
plan.iter()
.find(|p| p.layer == layer && p.canvas_id == canvas_id)
})
.map(|p| p.occlusions.clone())
.unwrap_or_default();
sticky_occlusions(
&self.frames[previous],
&self.frames[self.current_slot],
self.width,
self.height,
region,
s.delta.0,
s.delta.1,
&carry,
)
}
_ => Vec::new(),
};
if let Some(first) = occlusions.first() {
self.diagnostic(
Diagnostic::new("STICKY_OCCLUSION", Severity::Info)
.at(ctx.index, ctx.time)
.canvas(&canvas_id)
.region(*first),
);
}
let state = &mut self.states[i];
let placement = Placement {
layer,
canvas_id,
node: state
.last_node
.as_ref()
.map(|n| n.0.clone())
.unwrap_or_default(),
x: state.pose.0,
y: state.pose.1,
confidence: s.confidence,
uncertain: s.uncertain,
time: ctx.time,
skip: s.skip,
occlusions,
};
if moving {
state.blind = !s.textured;
state.weak = if s.decision == Decision::Tracked {
s.weak_step
} else {
false
};
}
state.previous_features = Some(s.own);
Outcome::Placed(placement, observed)
}
}