PWC-Net

PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume frames are extracted for every consecutive pair of frames in a video. PWC-Net is pre-trained on Sintel Flow dataset. The implementation follows sniklaus/pytorch-pwc@f61389005.

CUDA 11 and GPUs like RTX 3090 and newer

The current environment does not support CUDA 11 and, therefore, GPUs like RTX 3090 and newer. For details please check this issue #13 If you were able to fix it, please share your workaround. If you need an optical flow extractor, you are recommended to use RAFT.

The PWC-Net does NOT support using CPU currently

The PWC-Net uses cupy module, which makes it difficult to turn to a version that does not use the GPU. However, if you have solution, you may submit a PR.

Set up the Environment for PWC

Setup conda environment.

# it will create a new conda environment called 'pwc' on your machine
conda env create -f conda_env_pwc.yml

Quick Start

Activate the environment

conda activate pwc

and extract optical flow from ./sample/v_GGSY1Qvo990.mp4 and show the flow for each frame

python main.py \
    feature_type=pwc \
    device="cuda:0" \
    show_pred=true \
    video_paths="[./sample/v_GGSY1Qvo990.mp4]"

Note, if show_pred=true, the window with predictions will appear, use any key to show the next frame. To use show_pred=true, a screen must be attached to the machine or X11 forwarding is enabled.

Supported Arguments

Argument	Default	Description
`batch_size`	`1`	You may speed up extraction of features by increasing the batch size as much as your GPU permits.
`extraction_fps`	`null`	If specified (e.g. as `5`), the video will be re-encoded to the `extraction_fps` fps. Leave unspecified or `null` to skip re-encoding.
`side_size`	`null`	If resized to the smaller edge (`resize_to_smaller_edge=true`), then `min(W, H) = side_size`, if to the larger: max(W, H), if `null` (None) no resize is performed.
`resize_to_smaller_edge`	`true`	If `false`, the larger edge will be used to be resized to `side_size`.
`device`	`"cuda:0"`	The device specification. It follows the PyTorch style. Use `"cuda:3"` for the 4th GPU on the machine or `"cpu"` for CPU-only.
`video_paths`	`null`	A list of videos for feature extraction. E.g. `"[./sample/v_ZNVhz7ctTq0.mp4, ./sample/v_GGSY1Qvo990.mp4]"` or just one path `"./sample/v_GGSY1Qvo990.mp4"`.
`file_with_video_paths`	`null`	A path to a text file with video paths (one path per line). Hint: given a folder `./dataset` with `.mp4` files one could use: `find ./dataset -name "*mp4" > ./video_paths.txt`.
`on_extraction`	`print`	If `print`, the features are printed to the terminal. If `save_numpy` or `save_pickle`, the features are saved to either `.npy` file or `.pkl`.
`output_path`	`"./output"`	A path to a folder for storing the extracted features (if `on_extraction` is either `save_numpy` or `save_pickle`).
`keep_tmp_files`	`false`	If `true`, the reencoded videos will be kept in `tmp_path`.
`tmp_path`	`"./tmp"`	A path to a folder for storing temporal files (e.g. reencoded videos).
`show_pred`	`false`	If `true`, the script will visualize the optical flow for each pair of RGB frames.

Examples

Please see the examples for RAFT optical flow frame extraction. Make sure to replace --feature_type argument to pwc.

Credits

The PWC-Net paper and official implementation.
The PyTorch implementation used in this repo.

License

The wrapping code is under MIT, but PWC Net has GPL-3.0