Models

Builder

ezflow.models.build.build_model(name, cfg_path=None, custom_cfg=False, cfg=None, default=False, weights_path=None)[source]

Builds a model from a model name and config. Also supports loading weights

Parameters

  • name (str) – Name of the model to build

  • cfg_path (str, optional) – Path to a config file. If not provided, will use the default config for the model

  • custom_cfg (bool, optional) – Whether to use a custom config file. If False, will use the default config for the model

  • cfg (CfgNode object, optional) – Custom config object. If provided, will use this config instead of the default config for the model

  • default (bool, optional) – Whether to use the default config for the model

  • weights_path (str, optional) – Path to a weights file

Returns

The model

Return type

torch.nn.Module

DICL

class ezflow.models.dicl.DICL(cfg)[source]

Implementation of the paper Displacement-Invariant Matching Cost Learning for Accurate Optical Flow Estimation

Parameters

cfg (CfgNode) – Configuration for the model

forward(img1, img2)[source]

Performs forward pass of the network

Parameters

  • img1 (torch.Tensor) – Image to predict flow from

  • img2 (torch.Tensor) – Image to predict flow to

Returns

<flow_preds> torch.Tensor : intermediate flow predications from img1 to img2 <flow_upsampled> torch.Tensor : if model is in eval state, return upsampled flow

Return type

dict

FlowNetS

class ezflow.models.flownet_s.FlowNetS(cfg)[source]

Implementation of FlowNetSimple from the paper FlowNet: Learning Optical Flow with Convolutional Networks

Parameters

cfg (CfgNode) – Configuration for the model

forward(img1, img2)[source]

Performs forward pass of the network

Parameters

  • img1 (torch.Tensor) – Image to predict flow from

  • img2 (torch.Tensor) – Image to predict flow to

Returns

<flow_preds> torch.Tensor : intermediate flow predications from img1 to img2 <flow_upsampled> torch.Tensor : if model is in eval state, return upsampled flow

Return type

dict

FlowNetC

class ezflow.models.flownet_c.FlowNetC(cfg)[source]

Implementation of FlowNetCorrelation from the paper FlowNet: Learning Optical Flow with Convolutional Networks

Parameters

cfg (CfgNode) – Configuration for the model

forward(img1, img2)[source]

Performs forward pass of the network

Parameters

  • img1 (torch.Tensor) – Image to predict flow from

  • img2 (torch.Tensor) – Image to predict flow to

Returns

<flow_preds> torch.Tensor : intermediate flow predications from img1 to img2 <flow_upsampled> torch.Tensor : if model is in eval state, return upsampled flow

Return type

dict

Predictor

class ezflow.models.predictor.Predictor(model_name, model_cfg_path=None, model_cfg=None, model_weights_path=None, custom_cfg_file=False, default=False, data_transform=None, device='cpu', flow_scale=1.0, pad_divisor=1)[source]

A class that uses an instance of an optical flow estimation model to predict flow between two images

Parameters

  • model_name (str) – The name of the optical flow estimation model to use

  • model_cfg_path (str, optional) – The path to the config file for the optical flow estimation model, by default None in which case the default config is used

  • model_cfg (CfgNode object, optional) – The config object for the optical flow estimation model, by default None

  • model_weights_path (str, optional) – The path to the weights file for the optical flow estimation model

  • custom_cfg_file (bool, optional) – Whether the config file is a custom config file or one one of the configs included in EzFlow, by default False

  • default (bool, optional) – Whether to use the default config for the model

  • data_transform (torchvision.transforms object, optional) – The data transform to apply to the images before passing them to the model, by default None

  • device (str, optional) – The device to use for the model, by default “cpu”

  • flow_scale (float, optional) – The scale to apply to the predicted flow, by default 1.0

  • pad_divisor (int, optional) – The divisor to make the image dimensions evenly divisible by using padding, by default 1

PWCNet

class ezflow.models.pwcnet.PWCNet(cfg)[source]

Implementation of the paper PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume

Parameters

cfg (CfgNode) – Configuration for the model

forward(img1, img2)[source]

Performs forward pass of the network

Parameters

  • img1 (torch.Tensor) – Image to predict flow from

  • img2 (torch.Tensor) – Image to predict flow to

Returns

<flow_preds> torch.Tensor : intermediate flow predications from img1 to img2 <flow_upsampled> torch.Tensor : if model is in eval state, return upsampled flow

Return type

dict

RAFT

class ezflow.models.raft.RAFT(cfg)[source]

Implementation of the paper RAFT: Recurrent All-Pairs Field Transforms for Optical Flow

Parameters

cfg (CfgNode) – Configuration for the model

forward(img1, img2, flow_init=None)[source]

Performs forward pass of the network

Parameters

  • img1 (torch.Tensor) – Image to predict flow from

  • img2 (torch.Tensor) – Image to predict flow to

Returns

<flow_preds> torch.Tensor : intermediate flow predications from img1 to img2 <flow_upsampled> torch.Tensor : if model is in eval state, return upsampled flow

Return type

dict

VCN

class ezflow.models.vcn.VCN(cfg)[source]

Implementation of the paper Volumetric Correspondence Networks for Optical Flow

Parameters

cfg (CfgNode) – Configuration for the model

forward(img1, img2)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.