Color Fix for VapourSynth

For example for fixing color shift from AI upscaling/restoration models, or transfering color grading from an old release to a remaster. Also known as Color Transfer or Color Matching. See this collection of Comparisons.

Installation

Nvidia

pip install -U vs_colorfix[tensorrt] --extra-index-url https://pypi.nvidia.com/

To enable the Wavelet Color Fix CPU backend, install the plugin Vapoursynth-ATWT. (optional)

Others

pip install -U vs_colorfix

To enable the Wavelet Color Fix CPU backend, install the plugin Vapoursynth-ATWT. (optional)

[!TIP] For VapourSynth R73 and older, follow the manual installation steps.

Average Color Fix

Fixes color shift by matching the average color of a clip to a reference clip. A very fast way to transfer colors from one clip to another. For large color differences, the Wavelet Color Fix is more accurate.

import vs_colorfix
clip = vs_colorfix.average(clip, ref, radius=10, planes=[0, 1, 2], fast=False)

clip
Base clip where the colors will be applied to.
Recommended higher than 8-bit to avoid banding.

ref
Reference clip where the colors are taken from.
Check the comparisons to get an idea how close it should match the base clip.

radius
Higher means a more global color match and wider bloom/bleed.
Lower means a more local color match and smaller bloom/bleed. Too low and the reference clip will become visible.
Test values 5 and 30 and this will become more clear.

planes (optional)
Which planes to color fix. Any unmentioned planes will simply be copied.
If not set, all planes will be color-fixed.

fast (optional)
Does the averaging via a downscale instead of a blur, which is much faster, but will produce faint blocky artifacts.
I found it useful for radius > 60 where artifacts are no longer noticable, or to fix something like a prefilter clip.

[!TIP]

• If your clips are not sufficiently aligned or synchronized, use vs_align to align them first.
• To replicate chaiNNers Average Color Fix, you can convert % to radius: radius = (100/percentage-1)/2
  ChaiNNer works like fast=True does here, but it is recommended to leave it off for better results.

Wavelet Color Fix

Fixes color shift by first converting into wavelets, then matching the average color of a clip to a reference clip. Works similarly to the Average Color Fix, but more accurate when the color differences are large, at the cost of more computation.

import vs_colorfix
clip = vs_colorfix.wavelet(clip, ref, wavelets=4, planes=[0, 1, 2], backend="ncnn", num_streams=2, gpu_id=0, engine_folder=None)

clip
Base clip where the colors will be applied to.
Recommended higher than 8-bit to avoid banding.

ref
Reference clip where the colors are taken from.
Check the comparisons to get an idea how close it should match the base clip.

wavelets
Number of wavelets in the 1-10 range. Around 4 seems to work best in most cases.
Higher means a more global color match and wider bloom/bleed.
Lower means a more local color match and smaller bloom/bleed. Too low and the reference clip will become visible.
Test values 3 and 8 and this will become more clear.

planes (optional)
Which planes to color fix. Any unmentioned planes will simply be copied.
If not set, all planes will be color-fixed.

backend (optional)
The used backend. 16-bit float input is always much faster on GPU, but not supported by older GPUs.

• cpu CPU mode using the vapoursynth-ATWT plugin (slow).
• ncnn GPU mode using vs-mlrt with NCNN support. Works on almost any GPU, even Mac (fast).
• directml GPU mode using vs-mlrt with DirectML support. Works on most GPUs, Windows only (fast).
• tensorrt GPU mode using vs-mlrt with TensorRT support. Requires an Nvidia RTX GPU. On the first run, this mode will automatically build an engine, which may take a few minutes. Changing wavelets or input dimensions will trigger rebuilding, but build engines are stored (very fast).

num_streams (optional)
Number of parallel GPU streams. Higher can be faster, but requires more VRAM. Does not affect the CPU backend.

gpu_id (optional)
Which GPU to use starting from 0. Can be used to switch between iGPU/dGPU. Does not affect the CPU backend.

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

[!TIP]

• On low end GPUs, the NCNN backend can be faster than TensorRT and DirectML.
• If your clips are not sufficiently aligned or synchronized, use vs_align to align them first.

Benchmarks

Benchmarks were done on a RTX 4090 GPU and a Ryzen 5900X CPU with 16-bit input clips.

Acknowledgements

Average Color Fix idea from chaiNNer.
Wavelet Color Fix idea from sd-webui-stablesr.