Video Distortion Removal for VapourSynth

Also known as dewobble, warp stabilization, film or VHS distortion fix, atmospheric turbulence mitigation, or heat haze removal.

This is a partial implementation of the Turbulence Mitigation Transformer (only distortion removal, no deblurring). It does not do general video stabilization for shaky footage, only removes distortions within the frames. It is recommented to stabilize first if needed. Thanks to stinkybread for the comparison clips.

Installation

pip install -U vs_undistort --extra-index-url https://pypi.nvidia.com/

To enable the CPU/CUDA backends, install PyTorch with CUDA. (optional)

Usage

from vs_undistort import vs_undistort
clip = vs_undistort(clip, temp_window=10, window_overlap=0, interpolation="bicubic", backend="tensorrt", tiles=1, overlap=8, engine_folder=None)

clip
Distorted clip. Must be in RGBH format.

temp_window
Temporal window length. How many frames are grouped together and processed as a single chunk. Larger means higher VRAM requirements, but better temporal averaging and slower distortions can be removed. If this is too small, some distortions may not get removed, small jumps/hitches may be visible between windows and seams from tiling may become more obvious.

window_overlap (optional)
Overlap between temporal windows. Larger is slower. Increase if jumps/hitches between windows are noticeable.

interpolation (optional)
The interpolation mode used to warp the frames:

• bilinear More blurry.
• bicubic No blur, but may oversharpen slightly.

backend (optional)
The backend used to run the model:

• cpu CPU mode (very slow).
• cuda GPU mode using CUDA. Requires any Nvidia GPU (fast).
• tensorrt GPU mode using TensorRT. Requires an Nvidia RTX GPU. On the first run, this mode will automatically build an engine, which may take a few minutes. Changing interpolation, temp_window, or input dimensions will trigger rebuilding, but previously build engines are stored (very fast, low vram).

tiles (optional)
Amount of tiles to split the frames into. A higher amount reduces VRAM requirements, but also worsens spatial averaging. Default tiles=1 uses the full frame.

overlap (optional)
Overlap from one tile to the next. Increase if seams between tiles are visible. Does nothing when tiles=1.

engine_folder (optional)
Optional path to the TensorRT engine storage location. By default engines are stored in vs_undistort/engines. Only affects the TensorRT backend.

Benchmarks