Color enhancement aims to improve the visual quality of images and videos by restoring or enhancing color information. This process seeks to refine sharpness, saturation, and overall color balance, producing more natural and visually appealing results. To address this task, we'd like to propose a model that operates by first transforming RGB inputs into the YCbCr color space. The chrominance components (Cb and Cr) are individually denoised, while spatial feature extraction blocks and a global attention mechanism are employed to enhance the luminance component (Y). The model processes triplet frames to promote temporal consistency, and it is trained using a comprehensive set of loss functions, including perceptual loss, MS-SSIM loss, HSV color loss, edge loss, and optical flow-based temporal loss. Furthermore, the system is compatible with both PyTorch and ONNX frameworks and supports efficient multi-GPU inference pipelines, making it suitable for large-scale deployment.
- CUDA 11.8
- CUDNN 8.9
- Python 3.9
- Make Conda Environment
conda create -n ColEn python=3.9 -y
conda activate ColEn- Install Dependencies
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
conda install -c nvidia cudatoolkit=11.8 cudnn=8.9 -y
pip install -r requirements.txt
Run the following command for test:
- Pytorch model test:
python pytorch_video_test.py --video_path<input video path> --model_path<pytorch color-enhancement model> --save_dir<generated color enhance video path> - ONNX model test:
python onnx_video_test.py --video_path<input video path> --onnx_path<onnx color-enhancement model> ---save_dir<generated color enhance video path> | Input |  |
| Version-0326 |  |
| Input |  |
| Version-0326 |  |
| Input |  |
| Version-0326 |  |
| Input |  |
| Version-0326 |  |
| Input |  |
| Version-0326 |  |
Input: original
Output: enhanced
Input: original
Output: enhanced
Input: original
Output: enhanced