diff --git a/vidmap/frontend/keyframes/matching.py b/vidmap/frontend/keyframes/matching.py index 59f06de..53f3cf8 100644 --- a/vidmap/frontend/keyframes/matching.py +++ b/vidmap/frontend/keyframes/matching.py @@ -1,5 +1,7 @@ """Ordered low-resolution matching for keyframe selection.""" +from contextlib import contextmanager + import torch from torch.utils.data import DataLoader @@ -71,23 +73,51 @@ def build_pair_loader(scene_parser, sequence, lowres_options): return loader, len(pair_indices), original_width, original_height -def match_lowres_batch(tracker_model, batch, original_width, original_height, first_batch, *, batch_size): +@contextmanager +def pipelined_matches(tracker_model, loader, original_width, original_height, *, batch_size): + """Overlap one ordered inference batch with CPU selection, draining on exit.""" from vidmap.utils.profiling import record_timing, sync_time - names_A = batch["names_A"] - names_B = batch["names_B"] - batch_start = sync_time() - im_A = batch["im_A_batch"] - im_B = batch["im_B_batch"] - output = tracker_model.match_lowres_batch( - im_A, - im_B, - names_a=names_A, - names_b=names_B, - output_size=(original_width, original_height), - batch_size=batch_size, - ) - if first_batch: - record_timing("keyframing_first_batch", sync_time() - batch_start, first=True) - - return output.matches, output.certainty + source = iter(loader) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + first_batch = True + + def launch(batch): + nonlocal first_batch + with torch.cuda.stream(stream): + batch_start = sync_time() + output = tracker_model.match_lowres_batch( + batch["im_A_batch"], + batch["im_B_batch"], + names_a=batch["names_A"], + names_b=batch["names_B"], + output_size=(original_width, original_height), + batch_size=batch_size, + ) + if first_batch: + record_timing("keyframing_first_batch", sync_time() - batch_start, first=True) + ready = torch.cuda.Event() + ready.record() + first_batch = False + return (output.matches, output.certainty), ready + + def iterate(): + first = next(source, None) + pending = None if first is None else launch(first) + while pending is not None: + (matches, certainties), ready = pending + consumer = torch.cuda.current_stream() + consumer.wait_event(ready) + matches.record_stream(consumer) + certainties.record_stream(consumer) + following = next(source, None) + pending = None if following is None else launch(following) + yield matches, certainties + + batches = iterate() + try: + yield batches + finally: + stream.synchronize() + batches.close() diff --git a/vidmap/frontend/keyframes/processing.py b/vidmap/frontend/keyframes/processing.py index da36cd3..85e34dd 100644 --- a/vidmap/frontend/keyframes/processing.py +++ b/vidmap/frontend/keyframes/processing.py @@ -180,22 +180,21 @@ def select_candidates(self, salient_metadata) -> KeyframeCandidates: bootstrap_intrinsics, ) - first_batch = True - with tqdm( - total=total_pairs, - desc="Streaming keyframe detection", - disable=not progress_bars_enabled(), - ) as progress: - for batch in loader: - matches, certainties = keyframe_matching.match_lowres_batch( - tracker_model, - batch, - original_width, - original_height, - first_batch, - batch_size=self.lowres_options.batch_size, - ) - first_batch = False + with ( + keyframe_matching.pipelined_matches( + tracker_model, + loader, + original_width, + original_height, + batch_size=self.lowres_options.batch_size, + ) as batches, + tqdm( + total=total_pairs, + desc="Streaming keyframe detection", + disable=not progress_bars_enabled(), + ) as progress, + ): + for matches, certainties in batches: for pair_match_lr, pair_cert_lr in zip(matches, certainties): selector.process_pair(pair_match_lr, pair_cert_lr) progress.update(1) diff --git a/vidmap/frontend/options/matching.py b/vidmap/frontend/options/matching.py index 0db7a74..b0b6dfe 100644 --- a/vidmap/frontend/options/matching.py +++ b/vidmap/frontend/options/matching.py @@ -28,7 +28,7 @@ class PreprocessingOptions: @dataclass(frozen=True, config=ConfigDict(extra="forbid")) class LowresMatchOptions: batch_size: int = 8 - num_workers: int = 16 + num_workers: int = 4 resize_to_shape: Tuple[int, int] = (560, 560) interpolation: str = "torch_bicubic" diff --git a/vidmap/frontend/tracking/composition.py b/vidmap/frontend/tracking/composition.py index dec20d5..2e58abb 100644 --- a/vidmap/frontend/tracking/composition.py +++ b/vidmap/frontend/tracking/composition.py @@ -1,6 +1,7 @@ """Concrete composition of frontend's tracking-owned stages.""" import logging +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from vidmap.frontend.cache import fingerprint, ordered_files_fingerprint @@ -46,6 +47,10 @@ def __init__( def run(self): """Run model stages, then depth and optional camera calibration.""" + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="da3-verify") as preparation: + return self._run(preparation) + + def _run(self, preparation): from vidmap.frontend.pipeline import FrontendArtifacts, TrackingFrontendResult, validate_tracking_result options = self.options @@ -71,6 +76,7 @@ def run(self): logger.info("Geo-calibration (batch): %s", paths.geocalib_batch_path) frames = FrameSequence.from_scene(self.scene_parser) + verification = None with create_lazy_romav2_tracker(tracker_options) as tracker: keyframe_processor = KeyframeProcessor( scene_parser=self.scene_parser, @@ -116,6 +122,9 @@ def run(self): track_pairs_metadata=track_pairs_metadata, ) candidates = keyframe_processor.select_candidates(provisional_salient_metadata) + from vidmap.frontend.models.depth.da3_video import verify_da3_model_snapshot + + verification = preparation.submit(verify_da3_model_snapshot) keyframes, tracks = sparse_track_builder.admit_and_build_tracks( candidates, options.keyframes.selection if options.keyframes.selection.lookahead_pruning else None, @@ -151,6 +160,8 @@ def run(self): ) extended = extended_match_builder.build() + if verification is not None: + verification.result() depth_estimator = DepthEstimator( scene_parser=self.scene_parser, paths=paths,