README.md 4.1 KB
Newer Older
Max Ehrlich's avatar
Max Ehrlich committed
1
# Quantization Guided JPEG Artifact Correction
Max Ehrlich's avatar
Max Ehrlich committed
2

Max Ehrlich's avatar
Max Ehrlich committed
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
This repository contains the code and pretrained weights for the paper Quantization Guided JPEG Artifact Correction, published in ECCV 2020 and available on arXiv [here](https://arxiv.org/abs/2004.09320). 

> The JPEG image compression algorithm is the most popular method of image compression because of its ability for large compression ratios. However, to achieve such high compression, information is lost. For aggressive quantization settings, this leads to a noticeable reduction in image quality. Artifact correction has been studied in the context of deep neural networks for some time, but the current state-of-the-art methods require a different model to be trained for each quality setting, greatly limiting their practical application. We solve this problem by creating a novel architecture which is parameterized by the JPEG files quantization matrix. This allows our single model to achieve state-of-the-art performance over models trained for specific quality settings. 

For most use cases, you will not need the submodules of this repository, so there is no need to pass `--recurse-submodules` to `git clone`.

## Pre-trained Weights

The pre-trained weights are stored with the repository using [git large file storage](https://git-lfs.github.com/). Make sure that you have git-lfs installed before checking out the repository to ensure you get the weight files. This will also pull down required pre-computed DCT statistics and minc VGG weights for training with the GAN texture loss. The total size of the weights is appx 4.84GB.

## Correcting a Single JPEG Image

Use `experiments/correct_single.py` to correct a single JPEG file with the pretrained weights, the input file must be an actual JPEG file. You can use any weights as input, or merp between the color and GAN weights. 

example with color weights
```
19
python experiments/correct_single.py --input [filename].jpg --output restored.png --weights ./weights/weights_c.pt --dstats_path ./stats/cstats.pt --color
Max Ehrlich's avatar
Max Ehrlich committed
20 21
```

22
example with merp (70% GAN weights, 30% color weights)
Max Ehrlich's avatar
Max Ehrlich committed
23
```
24
python experiments/correct_single.py --input [filename].jpg --output restored.png --weights ./weights/weights_g.pt --dstats_path ./stats/cstats.pt --color --merp_from ./weights/weights_c.pt --merp_alpha 0.7
Max Ehrlich's avatar
Max Ehrlich committed
25 26
```

27 28 29 30
## Correcting a Folder of JPEG Images

Use `experiments/correct_batch.py` to correct a folder of JPEG images, it takes many of the same arguments as `correct_single.py`. The `--input` argument is a path to a folder of images with the `.jpg` extension, they must be real JPEG files. The `--output` argument should be a folder, if it does not exist it will be created. The images in the input folder will be corrected and then written to the output folder with the same name as PNG files (the extension will be changes from `.jpg` to `.png`). If your images are the same size, you can pass the `--batch-size` argument to correct multiple images at a time (as a "minibatch" of tensors). If your images are different sizes, leave `--batch-size` at its default size of 1. 

Max Ehrlich's avatar
Max Ehrlich committed
31 32
example with color weights, batch size of 1
```
Max Ehrlich's avatar
Max Ehrlich committed
33
python experiments/correct_batch.py --input [path] --output [path] --weights ./weights/weights_c.pt --dstats_path ./stats/cstats.pt --color
Max Ehrlich's avatar
Max Ehrlich committed
34 35
```

Max Ehrlich's avatar
Max Ehrlich committed
36 37 38 39 40 41 42
## Comparison Code

This repository references a repository called `quantization-guided-ac-comparisoncode`. This repository contains code from prior works which has small modifications to make easier to work with, allow for color images, fix bugs, etc. Since that code is not ours to redistribute (in general) that code is a private repository requiring permission to clone it. If you think you need access to it, please contact us with a good reason and we will make a decision on allowing you to access it.

## Acknowledgement

This project was partially supported by Facebook AI and Defense Advanced Research Projects Agency (DARPA) MediFor program (FA87501620191). The views, opinions, and findings expressed are those of the authors and should not be interpreted as representing the official views or policies of the Department of Defense or the U.S. Government. There is no collaboration between Facebook and DARPA.