mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-07 02:51:24 +00:00
* SAM 2 tracking: prompt frames in frame order, size-check frame reused (#3030) With several objects prompted on three or more frames, the prompts went to the tracker object by object. The tracker keeps features for two frames, so every prompt after the first object fetched and encoded its frame again. Each prompt frame was also fetched once only to check its size and then fetched again for the encoder. Prompts now go to the tracker in frame order (then object order), and the frame fetched for the size check is handed to the tracker the first time it asks for that frame (up to 256 MB of pixel data kept; beyond that the frame is fetched again as before). Objects do not influence each other, so the results are unchanged. Tests cover the prompt order, one host fetch per prompt frame, the order of several prompts of one object on one frame, and the memory cap. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> * SAM 2 tracking test: the held-frame cap is cumulative Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com> Co-authored-by: FredApps <153173854+FredApps@users.noreply.github.com>
This commit is contained in:
1 file changed
+135
-21
+135
-21
@@ -28,6 +28,35 @@ struct Source<'a> {
|
||||
count: usize,
|
||||
fetch: RefCell<Frames<'a>>,
|
||||
sizes: Rc<RefCell<HashMap<usize, (usize, usize)>>>,
|
||||
/// Prompt frames fetched only to learn their size, kept for the encoder (taken on first use).
|
||||
held: RefCell<Held>,
|
||||
}
|
||||
|
||||
/// Most pixel data (in bytes) kept from the size check of the prompt frames; frames beyond it are
|
||||
/// fetched again when the tracker needs them.
|
||||
const HELD_BYTES: usize = 256 << 20;
|
||||
|
||||
#[derive(Default)]
|
||||
struct Held {
|
||||
frames: HashMap<usize, Arc<ClipFrame>>,
|
||||
bytes: usize,
|
||||
}
|
||||
|
||||
impl Held {
|
||||
fn keep(&mut self, index: usize, f: &Arc<ClipFrame>) {
|
||||
let bytes = f.rgb.len().saturating_mul(std::mem::size_of::<[f32; 3]>());
|
||||
let fits = self.bytes.checked_add(bytes).filter(|t| *t <= HELD_BYTES);
|
||||
if let (Some(total), false) = (fits, self.frames.contains_key(&index)) {
|
||||
self.frames.insert(index, f.clone());
|
||||
self.bytes = total;
|
||||
}
|
||||
}
|
||||
|
||||
fn take(&mut self, index: usize) -> Option<Arc<ClipFrame>> {
|
||||
let f = self.frames.remove(&index)?;
|
||||
self.bytes = self.bytes.saturating_sub(f.rgb.len().saturating_mul(std::mem::size_of::<[f32; 3]>()));
|
||||
Some(f)
|
||||
}
|
||||
}
|
||||
|
||||
impl FrameSource for Source<'_> {
|
||||
@@ -41,7 +70,11 @@ impl FrameSource for Source<'_> {
|
||||
.filter(|_| index < self.count)
|
||||
.and_then(|offset| self.first.checked_add(offset))
|
||||
.ok_or_else(|| Error::Input(format!("frame index {index} is outside the clip")))?;
|
||||
let f = (self.fetch.borrow_mut())(number).ok_or_else(|| Error::Input(format!("frame {number} is not available")))?;
|
||||
let held = self.held.borrow_mut().take(index);
|
||||
let f = match held {
|
||||
Some(f) => f,
|
||||
None => (self.fetch.borrow_mut())(number).ok_or_else(|| Error::Input(format!("frame {number} is not available")))?,
|
||||
};
|
||||
self.sizes.borrow_mut().insert(index, (f.w, f.h));
|
||||
preprocess(&f.rgb, f.w, f.h)
|
||||
}
|
||||
@@ -58,6 +91,36 @@ fn to_frame(m: &FrameMask, size: (usize, usize)) -> Result<TrackedFrame> {
|
||||
/// Most objects one request may track.
|
||||
const MAX_OBJECTS: usize = 32;
|
||||
|
||||
/// Every prompt as `(object, frame index, prompt)`, in frame order (then object order), after
|
||||
/// checking each prompt frame and recording its size in `source`. Frame order keeps the tracker's
|
||||
/// two-frame feature cache useful: all objects prompted on a frame are encoded together, so each
|
||||
/// prompt frame is encoded once however many objects there are. Objects do not influence each
|
||||
/// other, so the order of prompts across objects does not change any result.
|
||||
fn gather_prompts<'r>(req: &'r TrackRequest, source: &Source<'_>) -> Result<Vec<(usize, usize, &'r effectcraft_segment::Prompt)>> {
|
||||
// Prompt coordinates are in the frame's pixels and the model works at 1024 x 1024, so the
|
||||
// prompted frames' sizes are needed first.
|
||||
let mut prompt_frames = vec![];
|
||||
for (oi, prompts) in req.objects.iter().enumerate() {
|
||||
for (frame, p) in prompts {
|
||||
let idx = frame
|
||||
.checked_sub(req.first)
|
||||
.and_then(|offset| usize::try_from(offset).ok())
|
||||
.filter(|i| *i < req.count)
|
||||
.ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?;
|
||||
if !source.sizes.borrow().contains_key(&idx) {
|
||||
let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?;
|
||||
validate_frame_size(f.w, f.h, f.rgb.len())?;
|
||||
source.sizes.borrow_mut().insert(idx, (f.w, f.h));
|
||||
source.held.borrow_mut().keep(idx, &f);
|
||||
}
|
||||
prompt_frames.push((oi, idx, p));
|
||||
}
|
||||
}
|
||||
// stable: prompts of one object on one frame keep their order
|
||||
prompt_frames.sort_by_key(|(oi, idx, _)| (*idx, *oi));
|
||||
Ok(prompt_frames)
|
||||
}
|
||||
|
||||
impl Sam2Mask {
|
||||
fn run_track(&self, req: &TrackRequest, frames: Frames<'_>, cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(usize, i64, TrackedFrame)) -> Result<()> {
|
||||
if req.count == 0 || req.count > MAX_CLIP {
|
||||
@@ -68,25 +131,8 @@ impl Sam2Mask {
|
||||
return Err(Error::Input(format!("tracking needs 1 to {MAX_OBJECTS} objects, each with a prompt on at least one frame")));
|
||||
}
|
||||
let sizes = Rc::new(RefCell::new(HashMap::new()));
|
||||
let source = Source { first: req.first, count: req.count, fetch: RefCell::new(frames), sizes: sizes.clone() };
|
||||
// Prompt coordinates are in the frame's pixels and the model works at 1024 x 1024, so the
|
||||
// prompted frames' sizes are needed first.
|
||||
let mut prompt_frames = vec![];
|
||||
for (oi, prompts) in req.objects.iter().enumerate() {
|
||||
for (frame, p) in prompts {
|
||||
let idx = frame
|
||||
.checked_sub(req.first)
|
||||
.and_then(|offset| usize::try_from(offset).ok())
|
||||
.filter(|i| *i < req.count)
|
||||
.ok_or_else(|| Error::Input(format!("prompt frame {frame} is outside the clip")))?;
|
||||
if !sizes.borrow().contains_key(&idx) {
|
||||
let f = (source.fetch.borrow_mut())(*frame).ok_or_else(|| Error::Input(format!("frame {frame} is not available")))?;
|
||||
validate_frame_size(f.w, f.h, f.rgb.len())?;
|
||||
sizes.borrow_mut().insert(idx, (f.w, f.h));
|
||||
}
|
||||
prompt_frames.push((oi, idx, p));
|
||||
}
|
||||
}
|
||||
let source = Source { first: req.first, count: req.count, fetch: RefCell::new(frames), sizes: sizes.clone(), held: RefCell::default() };
|
||||
let prompt_frames = gather_prompts(req, &source)?;
|
||||
let mut tracker = VideoTracker::new(self.video_model().clone(), Box::new(source))?;
|
||||
let object_id = |oi: usize| ObjectId(oi as u32 + 1);
|
||||
// each object's first prompted frame: its pass starts there, forwards and backwards
|
||||
@@ -174,8 +220,76 @@ mod tests {
|
||||
#[test]
|
||||
fn source_rejects_overflowing_and_out_of_range_indices_before_fetching() {
|
||||
let mut fetch = |_| panic!("invalid indices must not fetch a frame");
|
||||
let source = Source { first: i64::MAX, count: 2, fetch: RefCell::new(&mut fetch), sizes: Default::default() };
|
||||
let source = Source { first: i64::MAX, count: 2, fetch: RefCell::new(&mut fetch), sizes: Default::default(), held: RefCell::default() };
|
||||
assert!(matches!(source.frame(1), Err(Error::Input(_))));
|
||||
assert!(matches!(source.frame(usize::MAX), Err(Error::Input(_))));
|
||||
}
|
||||
|
||||
fn clip_frame(w: usize, h: usize) -> Arc<ClipFrame> {
|
||||
Arc::new(ClipFrame { rgb: vec![[0.5, 0.25, 0.75]; w * h], w, h })
|
||||
}
|
||||
|
||||
fn point_prompt() -> effectcraft_segment::Prompt {
|
||||
effectcraft_segment::Prompt { points: vec![([1.0, 1.0], true)], bbox: None, mask: None }
|
||||
}
|
||||
|
||||
/// Several objects prompted on the same frames: prompts come out frame by frame (so the
|
||||
/// two-frame feature cache holds each prompt frame for all of its objects), and each prompt
|
||||
/// frame is fetched from the host once, size check and encoding together.
|
||||
#[test]
|
||||
fn prompts_run_frame_major_and_each_prompt_frame_is_fetched_once() {
|
||||
let frames_prompted = [10i64, 13, 16];
|
||||
let req = TrackRequest { first: 10, count: 8, objects: (0..4).map(|_| frames_prompted.iter().map(|f| (*f, point_prompt())).collect()).collect() };
|
||||
let fetched = RefCell::new(vec![]);
|
||||
let mut fetch = |n: i64| {
|
||||
fetched.borrow_mut().push(n);
|
||||
Some(clip_frame(4, 3))
|
||||
};
|
||||
let source = Source { first: req.first, count: req.count, fetch: RefCell::new(&mut fetch), sizes: Default::default(), held: RefCell::default() };
|
||||
let order: Vec<(usize, usize)> = gather_prompts(&req, &source).unwrap().into_iter().map(|(oi, idx, _)| (idx, oi)).collect();
|
||||
let expected: Vec<(usize, usize)> = [0usize, 3, 6].iter().flat_map(|f| (0..4).map(move |o| (*f, o))).collect();
|
||||
assert_eq!(order, expected);
|
||||
assert_eq!(*fetched.borrow(), vec![10, 13, 16], "one size-check fetch per distinct prompt frame");
|
||||
// the tracker's first look at each prompt frame reuses the frame already fetched
|
||||
for idx in [0usize, 3, 6] {
|
||||
source.frame(idx).unwrap();
|
||||
}
|
||||
assert_eq!(fetched.borrow().len(), 3, "encoding a prompt frame must not fetch it again");
|
||||
// a held frame is handed over once; a second look fetches again
|
||||
source.frame(0).unwrap();
|
||||
assert_eq!(*fetched.borrow(), vec![10, 13, 16, 10]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompts_of_one_object_on_one_frame_keep_their_order() {
|
||||
let (mut a, mut b) = (point_prompt(), point_prompt());
|
||||
a.points = vec![([1.0, 1.0], true)];
|
||||
b.points = vec![([2.0, 2.0], false)];
|
||||
let req = TrackRequest { first: 0, count: 3, objects: vec![vec![(2, a.clone()), (1, point_prompt()), (2, b.clone())]] };
|
||||
let mut fetch = |_| Some(clip_frame(2, 2));
|
||||
let source = Source { first: 0, count: 3, fetch: RefCell::new(&mut fetch), sizes: Default::default(), held: RefCell::default() };
|
||||
let got = gather_prompts(&req, &source).unwrap();
|
||||
assert_eq!(got.iter().map(|(_, i, _)| *i).collect::<Vec<_>>(), vec![1, 2, 2]);
|
||||
assert_eq!(got[1].2.points, a.points);
|
||||
assert_eq!(got[2].2.points, b.points);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn held_frames_are_bounded() {
|
||||
let mut held = Held::default();
|
||||
let big = Arc::new(ClipFrame { rgb: vec![[0.0; 3]; HELD_BYTES / 12 + 1], w: 1, h: HELD_BYTES / 12 + 1 });
|
||||
held.keep(0, &big);
|
||||
assert!(held.take(0).is_none(), "a frame over the cap is not kept");
|
||||
let small = clip_frame(4, 4);
|
||||
held.keep(1, &small);
|
||||
assert!(held.take(1).is_some());
|
||||
assert_eq!(held.bytes, 0);
|
||||
// the cap is cumulative: frames that fit together are kept, the one that breaks it is not
|
||||
let third = Arc::new(ClipFrame { rgb: vec![[0.0; 3]; HELD_BYTES / 12 / 3 + 1], w: 1, h: HELD_BYTES / 12 / 3 + 1 });
|
||||
for i in 0..4 {
|
||||
held.keep(i, &third);
|
||||
}
|
||||
assert_eq!(held.frames.len(), 2);
|
||||
assert!(held.bytes <= HELD_BYTES);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user