Function compute_flows

flowmap/flow/__init__.py:24–33

(
    batch: Batch,
    flow_shape: tuple[int, int],
    device: torch.device,
    cfg: FlowPredictorCfg,
)

@torch.no_grad()
def compute_flows(
    batch: Batch,
    flow_shape: tuple[int, int],
    device: torch.device,
    cfg: FlowPredictorCfg,
) -> Flows:
    print("Precomputing optical flow.")
    flow_predictor = get_flow_predictor(cfg)
    flow_predictor.to(device)
    return flow_predictor.compute_bidirectional_flow(batch.to(device), flow_shape)