basicsr.archs.tof_arch

class basicsr.archs.tof_arch.BasicModule

Bases: Module

Basic module of SPyNet.

Note that unlike the architecture in spynet_arch.py, the basic module here contains batch normalization.

forward(tensor_input)

Parameters:

tensor_input (Tensor) – Input tensor with shape (b, 8, h, w). 8 channels contain: [reference image (3), neighbor image (3), initial flow (2)].

Returns:

Estimated flow with shape (b, 2, h, w)

Return type:

Tensor

training: bool

class basicsr.archs.tof_arch.SPyNetTOF(load_path=None)

Bases: Module

SPyNet architecture for TOF.

Note that this implementation is specifically for TOFlow. Please use spynet_arch.py for general use. They differ in the following aspects:

1. The basic modules here contain BatchNorm.

2. Normalization and denormalization are not done here, as they are done in TOFlow.

Paper: Optical Flow Estimation using a Spatial Pyramid Network

Reference: https://github.com/Coldog2333/pytoflow

Parameters:

load_path (str) – Path for pretrained SPyNet. Default: None.

forward(ref, supp)

Parameters:

• ref (Tensor) – Reference image with shape of (b, 3, h, w).

• supp – The supporting image to be warped: (b, 3, h, w).

Returns:

Estimated optical flow: (b, 2, h, w).

Return type:

Tensor

training: bool

class basicsr.archs.tof_arch.TOFlow(adapt_official_weights=False)

Bases: Module

PyTorch implementation of TOFlow.

In TOFlow, the LR frames are pre-upsampled and have the same size with the GT frames.

Paper: Video Enhancement with Task-Oriented Flow

Reference: https://github.com/anchen1011/toflow

Reference: https://github.com/Coldog2333/pytoflow

Parameters:

adapt_official_weights (bool) – Whether to adapt the weights translated from the official implementation. Set to false if you want to train from scratch. Default: False

denormalize(img)

forward(lrs)

Parameters:

lrs – Input lr frames: (b, 7, 3, h, w).

Returns:

SR frame: (b, 3, h, w).

Return type:

Tensor

normalize(img)

training: bool