├── .gitignore ├── LICENSE ├── README.md ├── evaluate.py ├── fr_model.py ├── models ├── fr_live_patchwise.model ├── fr_live_weighted.model ├── fr_tid_patchwise.model ├── fr_tid_weighted.model ├── nr_live_patchwise.model ├── nr_live_weighted.model ├── nr_tid_patchwise.model └── nr_tid_weighted.model └── nr_model.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | 27 | # PyInstaller 28 | # Usually these files are written by a python script from a template 29 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 30 | *.manifest 31 | *.spec 32 | 33 | # Installer logs 34 | pip-log.txt 35 | pip-delete-this-directory.txt 36 | 37 | # Unit test / coverage reports 38 | htmlcov/ 39 | .tox/ 40 | .coverage 41 | .coverage.* 42 | .cache 43 | nosetests.xml 44 | coverage.xml 45 | *,cover 46 | .hypothesis/ 47 | 48 | # Translations 49 | *.mo 50 | *.pot 51 | 52 | # Django stuff: 53 | *.log 54 | local_settings.py 55 | 56 | # Flask stuff: 57 | instance/ 58 | .webassets-cache 59 | 60 | # Scrapy stuff: 61 | .scrapy 62 | 63 | # Sphinx documentation 64 | docs/_build/ 65 | 66 | # PyBuilder 67 | target/ 68 | 69 | # IPython Notebook 70 | .ipynb_checkpoints 71 | 72 | # pyenv 73 | .python-version 74 | 75 | # celery beat schedule file 76 | celerybeat-schedule 77 | 78 | # dotenv 79 | .env 80 | 81 | # virtualenv 82 | venv/ 83 | ENV/ 84 | 85 | # Spyder project settings 86 | .spyderproject 87 | 88 | # Rope project settings 89 | .ropeproject 90 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2016 Dominique Maniry 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # deepIQA 2 | 3 | This is the reference implementation of [Deep Neural Networks for No-Reference and Full-Reference Image Quality Assessment][arxiv]. 4 | The pretrained models contained in the models directory were trained for both NR and FR IQA and for both model variants described in the paper. 5 | They were trained on the full LIVE or TID2013 database respectively, as used in the cross-dataset evaluations. This evaluation script uses non-overlapping 32x32 patches to produce deterministic scores, whereas the evaluation in the paper uses randomly sampled overlapping patches. 6 | 7 | > usage: evaluate.py [-h] [--model MODEL] [--top {patchwise,weighted}] 8 | > [--gpu GPU] 9 | > INPUT [REF] 10 | 11 | ## Dependencies 12 | * [chainer](http://chainer.org/) 13 | * ~~scikit-learn~~ 14 | * ~~opencv~~ 15 | 16 | ## TODO 17 | * add training code 18 | * add cpu support (minor change) 19 | * ~~remove opencv and scikit-learn dependencies for loading data (minor changes)~~ 20 | * ~~fix non-deterministic behaviour~~ 21 | 22 | [arxiv]: http://arxiv.org/abs/1612.01697 23 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python2 2 | import numpy as np 3 | from numpy.lib.stride_tricks import as_strided 4 | 5 | import chainer 6 | from chainer import computational_graph 7 | from chainer import cuda 8 | from chainer import optimizers 9 | from chainer import serializers 10 | 11 | import argparse 12 | import six 13 | import imageio 14 | import numbers 15 | 16 | from nr_model import Model 17 | from fr_model import FRModel 18 | 19 | 20 | def extract_patches(arr, patch_shape=(32,32,3), extraction_step=32): 21 | arr_ndim = arr.ndim 22 | 23 | if isinstance(patch_shape, numbers.Number): 24 | patch_shape = tuple([patch_shape] * arr_ndim) 25 | if isinstance(extraction_step, numbers.Number): 26 | extraction_step = tuple([extraction_step] * arr_ndim) 27 | 28 | patch_strides = arr.strides 29 | 30 | slices = tuple(slice(None, None, st) for st in extraction_step) 31 | indexing_strides = arr[slices].strides 32 | 33 | patch_indices_shape = ((np.array(arr.shape) - np.array(patch_shape)) // 34 | np.array(extraction_step)) + 1 35 | 36 | shape = tuple(list(patch_indices_shape) + list(patch_shape)) 37 | strides = tuple(list(indexing_strides) + list(patch_strides)) 38 | 39 | patches = as_strided(arr, shape=shape, strides=strides) 40 | return patches 41 | 42 | 43 | parser = argparse.ArgumentParser(description='evaluate.py') 44 | parser.add_argument('INPUT', help='path to input image') 45 | parser.add_argument('REF', default="", nargs="?", help='path to reference image, if omitted NR IQA is assumed') 46 | parser.add_argument('--model', '-m', default='', 47 | help='path to the trained model') 48 | parser.add_argument('--top', choices=('patchwise', 'weighted'), 49 | default='weighted', help='top layer and loss definition') 50 | parser.add_argument('--gpu', '-g', default=0, type=int, 51 | help='GPU ID') 52 | args = parser.parse_args() 53 | 54 | 55 | chainer.global_config.train = False 56 | chainer.global_config.cudnn_deterministic = True 57 | 58 | 59 | FR = True 60 | if args.REF == "": 61 | FR = False 62 | 63 | if FR: 64 | model = FRModel(top=args.top) 65 | else: 66 | model = Model(top=args.top) 67 | 68 | 69 | cuda.cudnn_enabled = True 70 | cuda.check_cuda_available() 71 | xp = cuda.cupy 72 | serializers.load_hdf5(args.model, model) 73 | model.to_gpu() 74 | 75 | 76 | if FR: 77 | ref_img = imageio.imread(args.REF) 78 | patches = extract_patches(ref_img) 79 | X_ref = np.transpose(patches.reshape((-1, 32, 32, 3)), (0, 3, 1, 2)) 80 | 81 | img = imageio.imread(args.INPUT) 82 | patches = extract_patches(img) 83 | X = np.transpose(patches.reshape((-1, 32, 32, 3)), (0, 3, 1, 2)) 84 | 85 | 86 | y = [] 87 | weights = [] 88 | batchsize = min(2000, X.shape[0]) 89 | t = xp.zeros((1, 1), np.float32) 90 | for i in six.moves.range(0, X.shape[0], batchsize): 91 | X_batch = X[i:i + batchsize] 92 | X_batch = xp.array(X_batch.astype(np.float32)) 93 | 94 | if FR: 95 | X_ref_batch = X_ref[i:i + batchsize] 96 | X_ref_batch = xp.array(X_ref_batch.astype(np.float32)) 97 | model.forward(X_batch, X_ref_batch, t, False, n_patches_per_image=X_batch.shape[0]) 98 | else: 99 | model.forward(X_batch, t, False, X_batch.shape[0]) 100 | 101 | y.append(xp.asnumpy(model.y[0].data).reshape((-1,))) 102 | weights.append(xp.asnumpy(model.a[0].data).reshape((-1,))) 103 | 104 | y = np.concatenate(y) 105 | weights = np.concatenate(weights) 106 | 107 | print("%f" % (np.sum(y*weights)/np.sum(weights))) 108 | -------------------------------------------------------------------------------- /fr_model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import chainer 4 | from chainer import Variable 5 | import chainer.functions as F 6 | import chainer.links as L 7 | from chainer import computational_graph 8 | from chainer import cuda 9 | from chainer import optimizers 10 | from chainer import serializers 11 | 12 | 13 | class FRModel(chainer.Chain): 14 | 15 | 16 | def __init__(self, top="patchwise"): 17 | super(FRModel, self).__init__( 18 | # feature extraction 19 | conv1 = L.Convolution2D(3, 32, 3, pad=1), 20 | conv2 = L.Convolution2D(32, 32, 3, pad=1), 21 | 22 | conv3 = L.Convolution2D(32, 64, 3, pad=1), 23 | conv4 = L.Convolution2D(64, 64, 3, pad=1), 24 | 25 | conv5 = L.Convolution2D(64, 128, 3, pad=1), 26 | conv6 = L.Convolution2D(128, 128, 3, pad=1), 27 | 28 | conv7 = L.Convolution2D(128, 256, 3, pad=1), 29 | conv8 = L.Convolution2D(256, 256, 3, pad=1), 30 | 31 | conv9 = L.Convolution2D(256, 512, 3, pad=1), 32 | conv10 = L.Convolution2D(512, 512, 3, pad=1), 33 | 34 | # quality regression 35 | fc1 = L.Linear(512 * 3, 512), 36 | fc2 = L.Linear(512, 1) 37 | 38 | ) 39 | 40 | self.top = top 41 | 42 | if top == "weighted": 43 | fc1_a = L.Linear(512 * 3, 512) 44 | fc2_a = L.Linear(512, 1) 45 | self.add_link("fc1_a", fc1_a) 46 | self.add_link("fc2_a", fc2_a) 47 | 48 | 49 | def extract_features(self, x, train=True): 50 | h = F.relu(self.conv1(x)) 51 | h = F.relu(self.conv2(h)) 52 | self.h1 = h 53 | h = F.max_pooling_2d(h,2) 54 | 55 | h = F.relu(self.conv3(h)) 56 | h = F.relu(self.conv4(h)) 57 | self.h2 = h 58 | h = F.max_pooling_2d(h,2) 59 | 60 | h = F.relu(self.conv5(h)) 61 | h = F.relu(self.conv6(h)) 62 | self.h3 = h 63 | h = F.max_pooling_2d(h,2) 64 | 65 | h = F.relu(self.conv7(h)) 66 | h = F.relu(self.conv8(h)) 67 | self.h4 = h 68 | h = F.max_pooling_2d(h,2) 69 | 70 | h = F.relu(self.conv9(h)) 71 | h = F.relu(self.conv10(h)) 72 | self.h5 = h 73 | h = F.max_pooling_2d(h,2) 74 | return h 75 | 76 | def forward(self, x_data, x_ref_data, y_data, train=True, 77 | n_patches_per_image=32): 78 | 79 | xp = cuda.cupy 80 | 81 | if not isinstance(x_data, Variable): 82 | x = Variable(x_data) 83 | else: 84 | x = x_data 85 | x_data = x.data 86 | 87 | self.n_images = y_data.shape[0] 88 | self.n_patches = x_data.shape[0] 89 | self.n_patches_per_image = n_patches_per_image 90 | x_ref = Variable(x_ref_data) 91 | 92 | h = self.extract_features(x) 93 | self.h = h 94 | 95 | h_ref = self.extract_features(x_ref) 96 | 97 | h = F.concat((h-h_ref, h, h_ref)) 98 | 99 | h_ = h # save intermediate features 100 | h = F.dropout(F.relu(self.fc1(h)), ratio=0.5) 101 | h = self.fc2(h) 102 | 103 | if self.top == "weighted": 104 | a = F.dropout(F.relu(self.fc1_a(h_)), ratio=0.5) 105 | a = F.relu(self.fc2_a(a))+0.000001 106 | t = Variable(y_data) 107 | self.weighted_loss(h, a, t) 108 | elif self.top == "patchwise": 109 | a = Variable(xp.ones_like(h.data)) 110 | t = Variable(xp.repeat(y_data, n_patches_per_image)) 111 | self.patchwise_loss(h, a, t) 112 | 113 | if train: 114 | return self.loss 115 | else: 116 | return self.loss, self.y 117 | 118 | def patchwise_loss(self, h, a, t): 119 | self.loss = F.sum(abs(h - F.reshape(t, (-1,1)))) 120 | self.loss /= self.n_patches 121 | if self.n_images > 1: 122 | h = F.split_axis(h, self.n_images, 0) 123 | a = F.split_axis(a, self.n_images, 0) 124 | else: 125 | h, a = [h], [a] 126 | self.y = h 127 | self.a = a 128 | 129 | def weighted_loss(self, h, a, t): 130 | self.loss = 0 131 | if self.n_images > 1: 132 | h = F.split_axis(h, self.n_images, 0) 133 | a = F.split_axis(a, self.n_images, 0) 134 | t = F.split_axis(t, self.n_images, 0) 135 | else: 136 | h, a, t = [h], [a], [t] 137 | 138 | for i in range(self.n_images): 139 | y = F.sum(h[i]*a[i], 0) / F.sum(a[i], 0) 140 | self.loss += abs(y - F.reshape(t[i], (1,))) 141 | self.loss /= self.n_images 142 | self.y = h 143 | self.a = a 144 | -------------------------------------------------------------------------------- /models/fr_live_patchwise.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/fr_live_patchwise.model -------------------------------------------------------------------------------- /models/fr_live_weighted.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/fr_live_weighted.model -------------------------------------------------------------------------------- /models/fr_tid_patchwise.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/fr_tid_patchwise.model -------------------------------------------------------------------------------- /models/fr_tid_weighted.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/fr_tid_weighted.model -------------------------------------------------------------------------------- /models/nr_live_patchwise.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/nr_live_patchwise.model -------------------------------------------------------------------------------- /models/nr_live_weighted.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/nr_live_weighted.model -------------------------------------------------------------------------------- /models/nr_tid_patchwise.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/nr_tid_patchwise.model -------------------------------------------------------------------------------- /models/nr_tid_weighted.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dmaniry/deepIQA/51dece994d73f43dbd2ea44d3f2ef581014194ad/models/nr_tid_weighted.model -------------------------------------------------------------------------------- /nr_model.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import chainer 4 | from chainer import Variable 5 | import chainer.functions as F 6 | import chainer.links as L 7 | from chainer import cuda 8 | from chainer import optimizers 9 | from chainer import serializers 10 | 11 | cuda.check_cuda_available() 12 | xp = cuda.cupy 13 | 14 | class Model(chainer.Chain): 15 | 16 | 17 | def __init__(self, top = "patchwise"): 18 | super(Model, self).__init__( 19 | conv1 = L.Convolution2D(3, 32, 3, pad=1), 20 | conv2 = L.Convolution2D(32, 32, 3, pad=1), 21 | 22 | conv3 = L.Convolution2D(32, 64, 3, pad=1), 23 | conv4 = L.Convolution2D(64, 64, 3, pad=1), 24 | 25 | conv5 = L.Convolution2D(64, 128, 3, pad=1), 26 | conv6 = L.Convolution2D(128, 128, 3, pad=1), 27 | 28 | conv7 = L.Convolution2D(128, 256, 3, pad=1), 29 | conv8 = L.Convolution2D(256, 256, 3, pad=1), 30 | 31 | conv9 = L.Convolution2D(256, 512, 3, pad=1), 32 | conv10 = L.Convolution2D(512, 512, 3, pad=1), 33 | 34 | fc1 = L.Linear(512, 512), 35 | fc2 = L.Linear(512, 1), 36 | 37 | fc1_a = L.Linear(512, 512), 38 | fc2_a = L.Linear(512, 1) 39 | ) 40 | 41 | self.top = top 42 | 43 | def forward(self, x_data, y_data, train=True, n_patches=32): 44 | 45 | if not isinstance(x_data, Variable): 46 | x = Variable(x_data) 47 | else: 48 | x = x_data 49 | x_data = x.data 50 | self.n_images = y_data.shape[0] 51 | self.n_patches = x_data.shape[0] 52 | self.n_patches_per_image = self.n_patches / self.n_images 53 | 54 | h = F.relu(self.conv1(x)) 55 | h = F.relu(self.conv2(h)) 56 | h = F.max_pooling_2d(h,2) 57 | 58 | h = F.relu(self.conv3(h)) 59 | h = F.relu(self.conv4(h)) 60 | h = F.max_pooling_2d(h,2) 61 | 62 | h = F.relu(self.conv5(h)) 63 | h = F.relu(self.conv6(h)) 64 | h = F.max_pooling_2d(h,2) 65 | 66 | h = F.relu(self.conv7(h)) 67 | h = F.relu(self.conv8(h)) 68 | h = F.max_pooling_2d(h,2) 69 | 70 | h = F.relu(self.conv9(h)) 71 | h = F.relu(self.conv10(h)) 72 | h = F.max_pooling_2d(h,2) 73 | 74 | h_ = h 75 | self.h = h_ 76 | 77 | h = F.dropout(F.relu(self.fc1(h_)), ratio=0.5) 78 | h = self.fc2(h) 79 | 80 | if self.top == "weighted": 81 | a = F.dropout(F.relu(self.fc1_a(h_)), ratio=0.5) 82 | a = F.relu(self.fc2_a(a))+0.000001 83 | t = Variable(y_data) 84 | self.weighted_loss(h, a, t) 85 | elif self.top == "patchwise": 86 | a = Variable(xp.ones_like(h.data)) 87 | t = Variable(xp.repeat(y_data, n_patches)) 88 | self.patchwise_loss(h, a, t) 89 | 90 | 91 | if train: 92 | return self.loss 93 | else: 94 | return self.loss, self.y 95 | 96 | 97 | def patchwise_loss(self, h, a, t): 98 | self.loss = F.sum(abs(h - F.reshape(t, (-1,1)))) 99 | self.loss /= self.n_patches 100 | if self.n_images > 1: 101 | h = F.split_axis(h, self.n_images, 0) 102 | a = F.split_axis(a, self.n_images, 0) 103 | else: 104 | h, a = [h], [a] 105 | self.y = h 106 | self.a = a 107 | 108 | def weighted_loss(self, h, a, t): 109 | self.loss = 0 110 | if self.n_images > 1: 111 | h = F.split_axis(h, self.n_images, 0) 112 | a = F.split_axis(a, self.n_images, 0) 113 | t = F.split_axis(t, self.n_images, 0) 114 | else: 115 | h, a, t = [h], [a], [t] 116 | 117 | for i in range(self.n_images): 118 | y = F.sum(h[i]*a[i], 0) / F.sum(a[i], 0) 119 | self.loss += abs(y - F.reshape(t[i], (1,))) 120 | self.loss /= self.n_images 121 | self.y = h 122 | self.a = a 123 | --------------------------------------------------------------------------------