Method __init__

flowmap/flow/flow_predictor_gmflow.py:28–64

(self, cfg: FlowPredictorGMFlowCfg)

class FlowPredictorGMFlow(FlowPredictor[FlowPredictorGMFlowCfg]):
    def __init__(self, cfg: FlowPredictorGMFlowCfg) -> None:
        super().__init__(cfg)

        # Warn that GMFlow isn't installed.
        if GMFlow is None:
            print(
                "Warning: GMFlow could not be imported. Did you forget to initialize "
                "the git submodules?"
            )
            sys.exit(1)

        # Ensure that the checkpoint exists.
        checkpoint = "gmflow-scale1-mixdata-train320x576-4c3a6e9a.pth"
        checkpoint_path = cfg.cache_path / checkpoint
        if not checkpoint_path.exists():
            checkpoint_path.parent.mkdir(exist_ok=True, parents=True)
            print("Downloading GMFlow checkpoint.")
            urllib.request.urlretrieve(
                f"https://s3.eu-central-1.amazonaws.com/avg-projects/unimatch/pretrained/{checkpoint}",
                checkpoint_path,
            )

        # Set up the model.
        self.model = GMFlow(
            feature_channels=128,
            num_scales=1,
            upsample_factor=8,
            num_head=1,
            attention_type="swin",
            ffn_dim_expansion=4,
            num_transformer_layers=6,
        )

        # Load the pre-trained checkpoint.
        checkpoint = torch.load(checkpoint_path)
        weights = checkpoint["model"] if "model" in checkpoint else checkpoint
        self.model.load_state_dict(weights, strict=False)

    def forward(
        self,