Function compute_tracks

flowmap/tracking/__init__.py:80–110

(
    batch: Batch,
    device: torch.device,
    tracking_cfg: TrackPredictorCfg,
    precomputation_cfg: TrackPrecomputationCfg,
)

def compute_tracks(
    batch: Batch,
    device: torch.device,
    tracking_cfg: TrackPredictorCfg,
    precomputation_cfg: TrackPrecomputationCfg,
) -> list[Tracks]:
    # Set up the tracker.
    tracker = get_track_predictor(tracking_cfg)
    tracker.to(device)

    # Since we only use tracks for overfitting, assert that the batch size is 1.
    b, _, _, _, _ = batch.videos.shape
    assert b == 1

    cache_key = get_cache_key(
        batch.datasets[0],
        batch.scenes[0],
        batch.indices[0],
        precomputation_cfg.interval,
        precomputation_cfg.radius,
    )
    disk_cache = make_cache(precomputation_cfg.cache_path)
    return disk_cache(
        cache_key,
        lambda: generate_video_tracks(
            tracker,
            batch.videos[:1].to(device),
            precomputation_cfg.interval,
            precomputation_cfg.radius,
        ),
    )