项目作者: bhar92

项目描述 :
Real-time video quality improvement for applications such as video-chat using Perceptual Losses
高级语言: Python
项目地址: git://github.com/bhar92/NeuralSuperResolution.git
创建时间: 2017-12-01T19:10:31Z
项目社区:https://github.com/bhar92/NeuralSuperResolution

开源协议:

下载


Real-time Super-Resolution for Video Quality Improvement

Table of Contents:

Introduction

This repository contains a pytorch implementation of an algorithm for single image super-resolution applied to frames from a web-cam. The model uses the method described in Perceptual Losses for Real-Time Style Transfer and Super-Resolution along with Instance Normalization. The code is based on the PyTorch Example code for fast-neural-style.

The program is written in Python, and uses the Pytorch library for generating and loading the CNN models.

Examples

An example of 4x super-resolution:
Shrek 4x Super-Resolution

An example of 8x super-resolution:
Shrek 8x Super-Resolution

Setup Instructions

All the programs were tested on the following setup:

Training Machine details:

  • CPU: 4x Intel Core i5-7300HQ CPU with 8GB RAM
  • GPU: Nvidia GeForce GTX 1050 Ti with 4GB VRAM
  • Operating System: Ubuntu 16.04.3 LTS
  • CUDA version: release 8.0, V8.0.61
  • CuDNN: 6.0.21
  • OpenCV: 3.3.0

Note: While a good GPU will help immensely with training the networks, it is not absolutely required to evaluate these programs.

Software Installation:

OpenCV Installation

  1. $ sudo apt-get install build-essential cmake git libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev python-dev python-numpy libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev libjasper-dev libdc1394-22-dev libavcodec-dev libavformat-dev libswscale-dev libv4l-dev libxvidcore-dev libx264-dev libgtk-3-dev libatlas-base-dev gfortran python2.7-dev python3.5-dev
Download OpenCV source:
  1. $ cd ~
  2. $ wget -O opencv.zip https://github.com/opencv/opencv/archive/3.3.0.zip
  3. $ unzip opencv.zip
  1. $ cd ~
  2. $ wget -O opencv_contrib.zip https://github.com/opencv/opencv_contrib/archive/3.3.0.zip
  3. $ unzip opencv_contrib.zip
Setup the Python virtualenv:
  1. Install the latest version of pip:
    1. $ sudo apt-get install python-pip && pip install --upgrade pip
  2. Install virtualenv and virtualenvwrapper:
    1. $ sudo pip install virtualenv virtualenvwrapper
    2. $ sudo rm -rf ~/.cache/pip
  3. Add the following lines to your ~/.bashrc file:
  1. export WORKON_HOME=$HOME/.virtualenvs
  2. source /usr/local/bin/virtualenvwrapper.sh
  1. Now source your ~/.bashrc:
    1. source ~/.bashrc
  2. Next create a virtual environment:
  1. mkvirtualenv supres -p python3
  1. To get back into the virtualenv, just type:
    1. workon supres
  2. Now, install numpy:
    1. pip install numpy
    Building OpenCV:
  3. Perform the cmake:
  1. $ cd ~/opencv-3.3.0/
  2. $ mkdir build
  3. $ cd build
  4. $ cmake -D CMAKE_BUILD_TYPE=RELEASE \
  5. -D CMAKE_INSTALL_PREFIX=/usr/local \
  6. -D INSTALL_PYTHON_EXAMPLES=ON \
  7. -D INSTALL_C_EXAMPLES=OFF \
  8. -D OPENCV_EXTRA_MODULES_PATH=~/opencv_contrib-3.3.0/modules \
  9. -D PYTHON_EXECUTABLE=~/.virtualenvs/cv/bin/python \
  10. -D BUILD_EXAMPLES=ON ..
  1. Build OpenCV!
    1. $ make -j4
  2. Now install OpenCV:
    1. $ sudo make install
    2. $ sudo ldconfig
  3. Some OpenCV bookkeeping:
    1. $ cd /usr/local/lib/python3.5/site-packages/
    2. $ ls -l
    Here, you should see a file like cv2.cpython-35m-x86_64-linux-gnu.so
    Go ahead and bravely rename it to cv2.so:
    1. sudo mv cv2.cpython-35m-x86_64-linux-gnu.so cv2.so
  4. Create a symlink in your virtual env:
    1. $ cd ~/.virtualenvs/supres/lib/python3.5/site-packages/
    2. $ ln -s /usr/local/lib/python3.5/site-packages/cv2.so cv2.so

    Install the other dependencies:

  5. Install pytorch
    1. pip3 install http://download.pytorch.org/whl/cu80/torch-0.2.0.post3-cp35-cp35m-manylinux1_x86_64.whl
    2. pip3 install torchvision
  6. Install the other requirements
    1. pip install -r requirements.txt

Usage

Running the live webcam demo

The code supports two scales of super-resolution: 4x and 8x.
Packaged with the code are two pretrained models:

  1. saved_models/coco4x_epoch_20.model for 4x Super-Resolution
  2. saved_models/coco8x_epoch_20.model for 8x Super-Resolution

In order to run the live demo for 4x Super-Resolution, run the following command:

  1. python super-resolution.py eval --model saved_models/coco4x_epoch_20.model --downsample-scale 4

In order to run the live demo for 8x Super-Resolution, run the following command:

  1. python super-resolution.py eval --model saved_models/coco8x_epoch_20.model --downsample-scale 8

These are all the options available in eval mode:

  1. $ python super-resolution.py eval --help
  2. usage: super-resolution.py eval [-h] --model MODEL --downsample-scale
  3. DOWNSAMPLE_SCALE
  4. optional arguments:
  5. -h, --help show this help message and exit
  6. --model MODEL saved model to be used for super resolution
  7. --downsample-scale DOWNSAMPLE_SCALE
  8. amount that you wish to downsample by. Default = 8

Training your own models

It is possible to train your own models with the super-resolution.py’s train mode.

The packaged models were trained with MS COCO dataset. In order to download the dataset and set aside 10k images for training, please run the download_dataset.sh script.

Once downloaded, you can run the following example command to train a model for 4x Super-Resolution:

  1. python super-resolution.py train --epochs 20 --dataset data/ --save-model-dir saved_models/ --checkpoint-model-dir checkpoints/ --downsample-scale 4

The full set of help options for train mode are as follows:

  1. $ python super-resolution.py train --help
  2. usage: super-resolution.py train [-h] [--epochs EPOCHS]
  3. [--batch-size BATCH_SIZE] [--dataset DATASET]
  4. --save-model-dir SAVE_MODEL_DIR
  5. [--checkpoint-model-dir CHECKPOINT_MODEL_DIR]
  6. [--lr LR] [--log-interval LOG_INTERVAL]
  7. [--checkpoint-interval CHECKPOINT_INTERVAL]
  8. [--downsample-scale DOWNSAMPLE_SCALE]
  9. optional arguments:
  10. -h, --help show this help message and exit
  11. --epochs EPOCHS number of training epochs, default is 2
  12. --batch-size BATCH_SIZE
  13. batch size for training, default is 4
  14. --dataset DATASET path to training dataset, the path should point to a
  15. folder containing another folder with all the training
  16. images
  17. --save-model-dir SAVE_MODEL_DIR
  18. path to folder where trained model will be saved.
  19. --checkpoint-model-dir CHECKPOINT_MODEL_DIR
  20. path to folder where checkpoints of trained models
  21. will be saved
  22. --lr LR learning rate, default is 1e-3
  23. --log-interval LOG_INTERVAL
  24. number of images after which the training loss is
  25. logged, default is 500
  26. --checkpoint-interval CHECKPOINT_INTERVAL
  27. number of batches after which a checkpoint of the
  28. trained model will be created
  29. --downsample-scale DOWNSAMPLE_SCALE
  30. amount that you wish to downsample by. Default = 8