Repository: photosynthesis-team/piq (language: python)

Issue:
Metric: IW-SSIM
**Is your feature request related to a problem? Please describe.**
Based on ordinary MS-SSIM this metric adds information weighting to assign different weights to different regions.

**Describe the solution you'd like**
Implementation can be based on MS-SSIM with additional weight maps

**Additional context**
Original paper: https://ece.uwaterloo.ca/~z70wang/publications/IWSSIM.pdf
MATLAB code: https://ece.uwaterloo.ca/~z70wang/research/iwssim/
Python code: https://github.com/Jack-guo-xy/Python-IW-SSIM


Metric: IW-SSIM
Closes #128 

## Proposed Changes

  - [x] Refactor SSIM to allow easy IW-SSIM implementation
  - [ ] IW-SSIM metric
  - [ ] IW-PSNR metric
  - [ ] Tests
  - [ ] Update README
MS-SSIM Tests. Bug in scale weights check.
Multi Scale SSIM uses scaling weights to compute a score for image pair. The [test 1](https://github.com/photosynthesis-team/piq/blob/19e3cb39b511cc6dbfd68ad7b564e3caa26a29b0/tests/test_ms_ssim.py#L78
) and [test 2](https://github.com/photosynthesis-team/piq/blob/19e3cb39b511cc6dbfd68ad7b564e3caa26a29b0/tests/test_ms_ssim.py#L201), which checks weight format, raises assertion error due to different purpose. 


**To Reproduce**
Steps to reproduce the behavior:
1. Run `pytest tests/test_ms_ssim.py` 
2. Observe no error

1. Open `test/test_ms_ssim.py`
2. Move to [line 78](https://github.com/photosynthesis-team/piq/blob/19e3cb39b511cc6dbfd68ad7b564e3caa26a29b0/tests/test_ms_ssim.py#L78)
3. See that AssetionError is raised due to missing `x` variable rather than wrong weight format.
4. Move to [line 201](https://github.com/photosynthesis-team/piq/blob/19e3cb39b511cc6dbfd68ad7b564e3caa26a29b0/tests/test_ms_ssim.py#L201)
5. See that AssetionError is raised due to missing `x` variable rather than wrong weight format.


**Suggestion**
- Fix the test
- Add assertion check, which is missing in [MS-SSIM](https://github.com/photosynthesis-team/piq/blob/master/piq/ms_ssim.py) function and class


Discussion / hints:

# [Codecov](https://codecov.io/gh/photosynthesis-team/piq/pull/201?src=pr&el=h1) Report
> Merging [#201](https://codecov.io/gh/photosynthesis-team/piq/pull/201?src=pr&el=desc) (f4ab9a3) into [master](https://codecov.io/gh/photosynthesis-team/piq/commit/7856cd079eed63111bb732ceb996321063e7753d?el=desc) (7856cd0) will **decrease** coverage by `0.15%`.
> The diff coverage is `80.00%`.

[![Impacted file tree graph](https://codecov.io/gh/photosynthesis-team/piq/pull/201/graphs/tree.svg?width=650&height=150&src=pr&token=8XS76R16J8)](https://codecov.io/gh/photosynthesis-team/piq/pull/201?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #201      +/-   ##
==========================================
- Coverage   95.33%   95.18%   -0.16%     
==========================================
  Files          26       26              
  Lines        1780     1785       +5     
==========================================
+ Hits         1697     1699       +2     
- Misses         83       86       +3     
```

| Flag | Coverage Δ | |
|---|---|---|
| unittests | `95.18% <80.00%> (-0.16%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/photosynthesis-team/piq/pull/201?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [piq/functional/filters.py](https://codecov.io/gh/photosynthesis-team/piq/pull/201/diff?src=pr&el=tree#diff-cGlxL2Z1bmN0aW9uYWwvZmlsdGVycy5weQ==) | `88.00% <40.00%> (-12.00%)` | :arrow_down: |
| [piq/ssim.py](https://codecov.io/gh/photosynthesis-team/piq/pull/201/diff?src=pr&el=tree#diff-cGlxL3NzaW0ucHk=) | `100.00% <100.00%> (ø)` | |

Kudos, SonarCloud Quality Gate passed!

[<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/bug.png' alt='Bug' width='16' height='16' />](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=BUG) [<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A.png' alt='A' width='16' height='16' />](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=BUG) [0 Bugs](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=BUG)  
[<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/vulnerability.png' alt='Vulnerability' width='16' height='16' />](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=VULNERABILITY) [<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A.png' alt='A' width='16' height='16' />](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=VULNERABILITY) [0 Vulnerabilities](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=VULNERABILITY)  
[<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/security_hotspot.png' alt='Security Hotspot' width='16' height='16' />](https://sonarcloud.io/project/security_hotspots?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=SECURITY_HOTSPOT) [<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A.png' alt='A' width='16' height='16' />](https://sonarcloud.io/project/security_hotspots?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=SECURITY_HOTSPOT) [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=SECURITY_HOTSPOT)  
[<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/code_smell.png' alt='Code Smell' width='16' height='16' />](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=CODE_SMELL) [<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A.png' alt='A' width='16' height='16' />](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=CODE_SMELL) [1 Code Smell](https://sonarcloud.io/project/issues?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&resolved=false&types=CODE_SMELL)

[<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/CoverageChart/NoCoverageInfo.png' alt='No Coverage information' width='16' height='16' />](https://sonarcloud.io/component_measures?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201) No Coverage information  
[<img src='https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/Duplications/3.png' alt='0.0%' width='16' height='16' />](https://sonarcloud.io/component_measures?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&metric=new_duplicated_lines_density&view=list) [0.0% Duplication](https://sonarcloud.io/component_measures?id=photosynthesis-team_photosynthesis.metrics&pullRequest=201&metric=new_duplicated_lines_density&view=list)


Closed temporally until further notice.


Fix patch:
diff --git a/README.rst b/README.rst
index 206bee09..cccd5095 100644
--- a/README.rst
+++ b/README.rst
@@ -170,6 +170,7 @@ Acronym      Year    Metric
 PSNR         \-      `Peak Signal-to-Noise Ratio <https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio>`_
 SSIM         2003    `Structural Similarity <https://en.wikipedia.org/wiki/Structural_similarity>`_
 MS-SSIM      2004    `Multi-Scale Structural Similarity <https://ieeexplore.ieee.org/abstract/document/1292216>`_
+IW-SSIM      2011    `Information Content Weighted Structural Similarity Index <https://ece.uwaterloo.ca/~z70wang/publications/IWSSIM.pdf>`_
 VIFp         2004    `Visual Information Fidelity <https://ieeexplore.ieee.org/document/1576816>`_
 FSIM         2011    `Feature Similarity Index Measure <https://ieeexplore.ieee.org/document/5705575>`_
 SR-SIM       2012    `Spectral Residual Based Similarity <https://sse.tongji.edu.cn/linzhang/ICIP12/ICIP-SR-SIM.pdf>`_
@@ -242,6 +243,7 @@ For metrics that can take greyscale or colour images, ``c`` means chromatic vers
 PSNR         0.687 / 0.496      0.687 / 0.496 `TID2013`_          0.676 / 0.488      \- / -                            0.407 / 0.276      0.407 / 0.277 `PIPAL`_
 SSIM         0.720 / 0.527      0.637 / 0.464 `TID2013`_          0.724 / 0.537      0.718 / 0.532 `KADID10k`_         0.498 / 0.345      0.529 / 0.369 `PIPAL`_
 MS-SSIM      0.798 / 0.597      0.787 / 0.608 `TID2013`_          0.802 / 0.609      0.802 / 0.609 `KADID10k`_         0.552 / 0.389      0.462 / -
+IW-SSIM      0.778 / 0.598      0.778 / 0.598 `IW-SSIM`_          0.850 / 0.666      0.850 / 0.666 `KADID10k`_         0.595 / 0.427      \- / -
 VIFp         0.610 / 0.458      0.610 / 0.457 `TID2013`_          0.650 / 0.477      0.650 / 0.477 `KADID10k`_         0.497 / 0.345      \- / -
 FSIM         0.802 / 0.629      0.801 / 0.630 `TID2013`_          0.830 / 0.639      0.829 / 0.639 `KADID10k`_         0.588 / 0.415      0.596 / 0.421 `PIPAL`_
 FSIMc        0.851 / 0.667      0.851 / 0.667 `TID2013`_          0.854 / 0.665      0.854 / 0.665 `KADID10k`_         0.590 / 0.416      \- / -
@@ -269,6 +271,7 @@ DISTS        0.805 / 0.613      0.830 / 0.639 `DISTS`_            0.875 / 0.695
 .. _DISTS: https://arxiv.org/abs/2004.07728
 .. _HaarPSI: https://arxiv.org/abs/1607.06140
 .. _PIPAL: https://arxiv.org/pdf/2011.15002.pdf
+.. _IW-SSIM: https://ieeexplore.ieee.org/document/7442122
 
 .. benchmark-section-end
 
diff --git a/docs/source/functions.rst b/docs/source/functions.rst
index a9a585fe..757e8531 100644
--- a/docs/source/functions.rst
+++ b/docs/source/functions.rst
@@ -15,6 +15,10 @@ Multi-Scale Structural Similarity (MS-SSIM)
 '''''''''''''''''''''''''''''''''''''''''''
 .. autofunction:: piq.multi_scale_ssim
 
+Information Content Weighted Structural Similarity (IW-SSIM)
+''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+.. autofunction:: piq.information_weighted_ssim
+
 Visual Information Fidelity (VIFp)
 ''''''''''''''''''''''''''''''''''
 .. autofunction:: piq.vif_p
diff --git a/docs/source/modules.rst b/docs/source/modules.rst
index 0ca178e1..eaedc616 100644
--- a/docs/source/modules.rst
+++ b/docs/source/modules.rst
@@ -13,6 +13,11 @@ Multi-Scale Structural Similarity (MS-SSIM)
 .. autoclass:: piq.MultiScaleSSIMLoss
     :members:
 
+Information Content Weighted Structural Similarity (IW-SSIM)
+''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+.. autoclass:: piq.InformationWeightedSSIMLoss
+    :members:
+
 Visual Information Fidelity (VIFp)
 ''''''''''''''''''''''''''''''''''
 .. autoclass:: piq.VIFLoss
diff --git a/piq/__init__.py b/piq/__init__.py
index 63e51d93..b0471219 100644
--- a/piq/__init__.py
+++ b/piq/__init__.py
@@ -21,3 +21,4 @@
 from .srsim import srsim, SRSIMLoss
 from .pieapp import PieAPP
 from .dss import dss, DSSLoss
+from .iw_ssim import information_weighted_ssim, InformationWeightedSSIMLoss
diff --git a/piq/functional/__init__.py b/piq/functional/__init__.py
index 734dfd5c..dee6c959 100644
--- a/piq/functional/__init__.py
+++ b/piq/functional/__init__.py
@@ -1,6 +1,7 @@
 from piq.functional.base import ifftshift, get_meshgrid, similarity_map, gradient_map, pow_for_complex, crop_patches
 from piq.functional.colour_conversion import rgb2lmn, rgb2xyz, xyz2lab, rgb2lab, rgb2yiq, rgb2lhm
 from piq.functional.filters import haar_filter, hann_filter, scharr_filter, prewitt_filter, gaussian_filter
+from piq.functional.filters import binomial_filter1d, average_filter2d
 from piq.functional.layers import L2Pool2d
 from piq.functional.resize import imresize
 
@@ -8,6 +9,7 @@
     'ifftshift', 'get_meshgrid', 'similarity_map', 'gradient_map', 'pow_for_complex', 'crop_patches',
     'rgb2lmn', 'rgb2xyz', 'xyz2lab', 'rgb2lab', 'rgb2yiq', 'rgb2lhm',
     'haar_filter', 'hann_filter', 'scharr_filter', 'prewitt_filter', 'gaussian_filter',
+    'binomial_filter1d', 'average_filter2d',
     'L2Pool2d',
     'imresize',
 ]
diff --git a/piq/functional/filters.py b/piq/functional/filters.py
index ec5dedf7..69b508c0 100644
--- a/piq/functional/filters.py
+++ b/piq/functional/filters.py
@@ -1,5 +1,6 @@
 r"""Filters for gradient computation, bluring, etc."""
 import torch
+import numpy as np
 
 
 def haar_filter(kernel_size: int) -> torch.Tensor:
@@ -56,3 +57,30 @@ def prewitt_filter() -> torch.Tensor:
     Returns:
         kernel: Tensor with shape (1, 3, 3)"""
     return torch.tensor([[[-1., 0., 1.], [-1., 0., 1.], [-1., 0., 1.]]]) / 3
+
+
+def binomial_filter1d(kernel_size: int) -> torch.Tensor:
+    r"""Creates 1D normalized binomial filter
+
+    Args:
+        kernel_size (int): kernel size
+
+    Returns:
+        Binomial kernel with shape (1, 1, kernel_size)
+    """
+    kernel = np.poly1d([0.5, 0.5]) ** (kernel_size - 1)
+    return torch.tensor(kernel.c).view(1, 1, kernel_size)
+
+
+def average_filter2d(kernel_size: int) -> torch.Tensor:
+    r"""Creates 2D normalized average filter
+
+    Args:
+        kernel_size (int):
+
+    Returns:
+        kernel: Tensor with shape (1, kernel_size, kernel_size)
+    """
+    window = torch.ones(kernel_size) / kernel_size
+    kernel = window[:, None] * window[None, :]
+    return kernel.unsqueeze(0)
diff --git a/piq/iw_ssim.py b/piq/iw_ssim.py
new file mode 100644
index 00000000..cc55d3d7
--- /dev/null
+++ b/piq/iw_ssim.py
@@ -0,0 +1,457 @@
+r""" Implementation of Information Content Weighted Structural Similarity Index (IW-SSIM).
+
+Information Content Weighted Structural Similarity Index (IW-SSIM) [1] is an extension of
+the structural similarity (SSIM). IW-SSIM uses the idea of information content weighted pooling for similarity
+evaluation.
+
+Estimation values produced by presented implementation corresponds to MATLAB based estimations [2].
+
+References:
+  [1] Wang, Zhou, and Qiang Li.
+    "Information content weighting for perceptual image quality assessment."
+    IEEE Transactions on image processing 20.5 (2011): 1185-1198.
+    https://ece.uwaterloo.ca/~z70wang/publications/IWSSIM.pdf
+
+  [2] https://ece.uwaterloo.ca/~z70wang/research/iwssim/
+"""
+
+import torch
+from torch.nn.modules.loss import _Loss
+import torch.nn.functional as F
+from piq.utils import _validate_input, _reduce, _parse_version
+from piq.functional import gaussian_filter, binomial_filter1d, average_filter2d, rgb2yiq
+from typing import Union, Optional, Tuple
+import math
+
+
+def information_weighted_ssim(x: torch.Tensor, y: torch.Tensor, data_range: Union[int, float] = 1.,
+                              kernel_size: int = 11, kernel_sigma: float = 1.5, k1: float = 0.01, k2: float = 0.03,
+                              parent: bool = True, blk_size: int = 3, sigma_nsq: float = 0.4,
+                              scale_weights: Optional[torch.Tensor] = None,
+                              reduction: str = 'mean') -> torch.Tensor:
+    r"""Interface of Information Content Weighted Structural Similarity (IW-SSIM) index.
+    Inputs supposed to be in range ``[0, data_range]``.
+
+    Args:
+        x: An input tensor. Shape :math:`(N, C, H, W)`.
+        y: A target tensor. Shape :math:`(N, C, H, W)`.
+        data_range: Maximum value range of images (usually 1.0 or 255).
+        kernel_size: The side-length of the sliding window used in comparison. Must be an odd value.
+        kernel_sigma: Sigma of normal distribution for sliding window used in comparison.
+        k1: Algorithm parameter, K1 (small constant).
+        k2: Algorithm parameter, K2 (small constant).
+            Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results.
+        parent: Flag to control dependency on previous layer of pyramid.
+        blk_size: The side-length of the sliding window used in comparison for information content.
+        sigma_nsq: Parameter of visual distortion model.
+        scale_weights: Weights for scaling.
+        reduction: Specifies the reduction type:
+            ``'none'`` | ``'mean'`` | ``'sum'``. Default:``'mean'``
+
+    Returns:
+        Value of Information Content Weighted Structural Similarity (IW-SSIM) index.
+
+    References:
+        Wang, Zhou, and Qiang Li..
+        Information content weighting for perceptual image quality assessment.
+        IEEE Transactions on image processing 20.5 (2011): 1185-1198.
+        https://ece.uwaterloo.ca/~z70wang/publications/IWSSIM.pdf DOI:`10.1109/TIP.2010.2092435`
+
+    Note:
+        Lack of content in target image could lead to RuntimeError due to singular information content matrix,
+        which cannot be inverted.
+    """
+    assert kernel_size % 2 == 1, f'Kernel size must be odd, got [{kernel_size}]'
+
+    _validate_input(tensors=[x, y], dim_range=(4, 4), data_range=(0., data_range))
+
+    x = x / float(data_range) * 255
+    y = y / float(data_range) * 255
+
+    if x.size(1) == 3:
+        x = rgb2yiq(x)[:, :1]
+        y = rgb2yiq(y)[:, :1]
+
+    if scale_weights is None:
+        scale_weights = torch.tensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333], dtype=x.dtype, device=x.device)
+    scale_weights = scale_weights / scale_weights.sum()
+    if scale_weights.size(0) != scale_weights.numel():
+        raise ValueError(f'Expected a vector of weights, got {scale_weights.dim()}D tensor')
+
+    levels = scale_weights.size(0)
+
+    min_size = (kernel_size - 1) * 2 ** (levels - 1) + 1
+    if x.size(-1) < min_size or x.size(-2) < min_size:
+        raise ValueError(f'Invalid size of the input images, expected at least {min_size}x{min_size}.')
+
+    blur_pad = math.ceil((kernel_size - 1) / 2)  # Ceil
+    iw_pad = blur_pad - math.floor((blk_size - 1) / 2)  # floor
+    gauss_kernel = gaussian_filter(kernel_size, kernel_sigma).repeat(x.size(1), 1, 1, 1).to(x)
+
+    # Size of the kernel size to build Laplacian pyramid
+    pyramid_kernel_size = 5
+    bin_filter = binomial_filter1d(kernel_size=pyramid_kernel_size).to(x) * 2 ** 0.5
+
+    lo_x, x_diff_old = _pyr_step(x, bin_filter)
+    lo_y, y_diff_old = _pyr_step(y, bin_filter)
+
+    x = lo_x
+    y = lo_y
+    wmcs = []
+
+    for i in range(levels):
+        if i < levels - 2:
+            lo_x, x_diff = _pyr_step(x, bin_filter)
+            lo_y, y_diff = _pyr_step(y, bin_filter)
+            x = lo_x
+            y = lo_y
+
+        else:
+            x_diff = x
+            y_diff = y
+
+        ssim_map, cs_map = _ssim_per_channel(x=x_diff_old, y=y_diff_old, kernel=gauss_kernel, data_range=255,
+                                             k1=k1, k2=k2)
+
+        if parent and i < levels - 2:
+            iw_map = _information_content(x=x_diff_old, y=y_diff_old, y_parent=y_diff, kernel_size=blk_size,
+                                          sigma_nsq=sigma_nsq)
+
+            iw_map = iw_map[:, :, iw_pad:-iw_pad, iw_pad:-iw_pad]
+
+        elif i == levels - 1:
+            iw_map = torch.ones_like(cs_map)
+            cs_map = ssim_map
+
+        else:
+            iw_map = _information_content(x=x_diff_old, y=y_diff_old, y_parent=None, kernel_size=blk_size,
+                                          sigma_nsq=sigma_nsq)
+            iw_map = iw_map[:, :, iw_pad:-iw_pad, iw_pad:-iw_pad]
+
+        wmcs.append(torch.sum(cs_map * iw_map, dim=(-2, -1)) / torch.sum(iw_map, dim=(-2, -1)))
+
+        x_diff_old = x_diff
+        y_diff_old = y_diff
+
+    wmcs = torch.stack(wmcs, dim=0).abs()
+
+    score = torch.prod((wmcs ** scale_weights.view(-1, 1, 1)), dim=0)[:, 0]
+
+    return _reduce(x=score, reduction=reduction)
+
+
+class InformationWeightedSSIMLoss(_Loss):
+    r"""Creates a criterion that measures the Interface of Information Content Weighted Structural Similarity (IW-SSIM)
+    index error betweeneach element in the input :math:`x` and target :math:`y`.
+
+    Inputs supposed to be in range ``[0, data_range]``.
+
+    If :attr:`reduction` is not ``'none'`` (default ``'mean'``), then:
+
+    .. math::
+        InformationWeightedSSIMLoss(x, y) =
+        \begin{cases}
+            \operatorname{mean}(1 - IWSSIM), &  \text{if reduction} = \text{'mean';}\\
+            \operatorname{sum}(1 - IWSSIM),  &  \text{if reduction} = \text{'sum'.}
+        \end{cases}
+
+    Args:
+        data_range: Maximum value range of images (usually 1.0 or 255).
+        kernel_size: The side-length of the sliding window used in comparison. Must be an odd value.
+        kernel_sigma: Sigma of normal distribution for sliding window used in comparison.
+        k1: Algorithm parameter, K1 (small constant).
+        k2: Algorithm parameter, K2 (small constant).
+            Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results.
+        parent: Flag to control dependency on previous layer of pyramid.
+        blk_size: The side-length of the sliding window used in comparison for information content.
+        sigma_nsq: Sigma of normal distribution for sliding window used in comparison for information content.
+        scale_weights: Weights for scaling.
+        reduction: Specifies the reduction type:
+            ``'none'`` | ``'mean'`` | ``'sum'``. Default:``'mean'``
+
+    Examples:
+        >>> loss = InformationWeightedSSIMLoss()
+        >>> input = torch.rand(3, 3, 256, 256, requires_grad=True)
+        >>> target = torch.rand(3, 3, 256, 256)
+        >>> output = loss(input, target)
+        >>> output.backward()
+
+    References:
+        Wang, Zhou, and Qiang Li..
+        Information content weighting for perceptual image quality assessment.
+        IEEE Transactions on image processing 20.5 (2011): 1185-1198.
+        https://ece.uwaterloo.ca/~z70wang/publications/IWSSIM.pdf DOI:`10.1109/TIP.2010.2092435`
+
+    """
+
+    def __init__(self, data_range: Union[int, float] = 1., kernel_size: int = 11, kernel_sigma: float = 1.5,
+                 k1: float = 0.01, k2: float = 0.03, parent: bool = True, blk_size: int = 3, sigma_nsq: float = 0.4,
+                 scale_weights: Optional[torch.Tensor] = None, reduction: str = 'mean'):
+        super().__init__()
+        self.data_range = data_range
+        self.kernel_size = kernel_size
+        self.kernel_sigma = kernel_sigma
+        self.k1 = k1
+        self.k2 = k2
+        self.parent = parent
+        self.blk_size = blk_size
+        self.sigma_nsq = sigma_nsq
+        self.scale_weights = scale_weights
+        self.reduction = reduction
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        r"""Computation of Information Content Weighted Structural Similarity (IW-SSIM) index as a loss function.
+        For colour images channel order is RGB.
+
+        Args:
+            x: An input tensor. Shape :math:`(N, C, H, W)`.
+            y: A target tensor. Shape :math:`(N, C, H, W)`.
+
+        Returns:
+            Value of IW-SSIM loss to be minimized, i.e. ``1 - information_weighted_ssim`` in [0, 1] range.
+        """
+        score = information_weighted_ssim(x=x, y=y, data_range=self.data_range, kernel_size=self.kernel_size,
+                                          kernel_sigma=self.kernel_sigma, k1=self.k1, k2=self.k2,
+                                          parent=self.parent, blk_size=self.blk_size, sigma_nsq=self.sigma_nsq,
+                                          scale_weights=self.scale_weights, reduction=self.reduction)
+        return torch.ones_like(score) - score
+
+
+def _pyr_step(x: torch.Tensor, kernel: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    r""" Computes one step of Laplacian pyramid generation.
+
+    Args:
+        x: An input tensor. Shape :math:`(N, C, H, W)`.
+        kernel: Kernel to perform blurring.
+
+    Returns:
+        Tuple of tensors with downscaled low resolution image and high-resolution difference.
+    """
+    # Blur and Downsampling
+    up_pad = (kernel.size(-1) - 1) // 2  # 5 -> 2, 4 -> 1
+    down_pad = kernel.size(-1) - 1 - up_pad  # 5 -> 2, 4 -> 2
+    kernel_t = kernel.transpose(-2, -1)
+    lo_x = x
+    if x.size(-1) > 1:
+        lo_x = F.pad(lo_x, pad=[up_pad, down_pad, 0, 0], mode='reflect')
+        lo_x = F.conv2d(input=lo_x, weight=kernel.unsqueeze(0), padding=0)[:, :, :, ::2]
+    if x.size(-2) > 1:
+        lo_x = F.pad(lo_x, pad=[0, 0, up_pad, down_pad], mode='reflect')
+        lo_x = F.conv2d(input=lo_x, weight=kernel_t.unsqueeze(0), padding=0)[:, :, ::2, :]
+
+    # Upsampling and Blur
+    up_pad = (kernel.size(-1) - 1) // 2  # 5 -> 2, 4 -> 1
+    down_pad = kernel.size(-1) - 1 - up_pad  # 5 -> 2, 4 -> 2
+    hi_x = lo_x
+
+    if x.size(-1) > 1:
+        upsampling_kernel = torch.tensor([[[[1., 0.]]]], dtype=x.dtype, device=x.device)
+        hi_x = F.conv_transpose2d(input=hi_x, weight=upsampling_kernel, stride=(1, 2), padding=0)
+        hi_x = F.pad(hi_x, pad=[up_pad, down_pad, 0, 0], mode='reflect')
+        hi_x = F.conv2d(input=hi_x, weight=kernel.unsqueeze(0), padding=0)[:, :, :, :x.size(-1)]
+    if x.size(-2) > 1:
+        upsampling_kernel = torch.tensor([[[[1.], [0.]]]], dtype=x.dtype, device=x.device)
+        hi_x = F.conv_transpose2d(input=hi_x, weight=upsampling_kernel, stride=(2, 1), padding=0)
+        hi_x = F.pad(hi_x, pad=[0, 0, up_pad, down_pad], mode='reflect')
+        hi_x = F.conv2d(input=hi_x, weight=kernel_t.unsqueeze(0), padding=0)[:, :, :x.size(-2), :]
+
+    hi_x = x - hi_x
+    return lo_x, hi_x
+
+
+def _ssim_per_channel(x: torch.Tensor, y: torch.Tensor, kernel: torch.Tensor,
+                      data_range: Union[float, int] = 1., k1: float = 0.01,
+                      k2: float = 0.03) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""Computes Structural Similarity (SSIM) index for X and Y per channel.
+
+    Args:
+        x: An input tensor. Shape :math:`(N, C, H, W)`.
+        y: A target tensor. Shape :math:`(N, C, H, W)`.
+        kernel: 2D Gaussian kernel.
+        data_range: Maximum value range of images (usually 1.0 or 255).
+        k1: Algorithm parameter, K1 (small constant, see [1]).
+        k2: Algorithm parameter, K2 (small constant, see [1]).
+            Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results.
+
+    Returns:
+        Tuple with Structural Similarity maps and Contrast maps.
+    """
+    c1 = (k1 * data_range) ** 2
+    c2 = (k2 * data_range) ** 2
+    n_channels = x.size(1)
+    mu_x = F.conv2d(x, weight=kernel, stride=1, padding=0, groups=n_channels)
+    mu_y = F.conv2d(y, weight=kernel, stride=1, padding=0, groups=n_channels)
+
+    mu_xx = mu_x ** 2
+    mu_yy = mu_y ** 2
+    mu_xy = mu_x * mu_y
+
+    sigma_xx = F.conv2d(x ** 2, weight=kernel, stride=1, padding=0, groups=n_channels) - mu_xx
+    sigma_yy = F.conv2d(y ** 2, weight=kernel, stride=1, padding=0, groups=n_channels) - mu_yy
+    sigma_xy = F.conv2d(x * y, weight=kernel, stride=1, padding=0, groups=n_channels) - mu_xy
+
+    sigma_xx = F.relu(sigma_xx)
+    sigma_yy = F.relu(sigma_yy)
+
+    # Contrast sensitivity (CS) with alpha = beta = gamma = 1.
+    cs_map = (2. * sigma_xy + c2) / (sigma_xx + sigma_yy + c2)
+
+    # Structural similarity (SSIM)
+    ss_map = (2. * mu_xy + c1) / (mu_xx + mu_yy + c1) * cs_map
+
+    return ss_map, cs_map
+
+
+def _information_content(x: torch.Tensor, y: torch.Tensor, y_parent: torch.Tensor = None,
+                         kernel_size: int = 3, sigma_nsq: float = 0.4) -> torch.Tensor:
+    r"""Computes Information Content Map for weighting the Structural Similarity.
+
+    Args:
+        x: An input tensor. Shape :math:`(N, C, H, W)`.
+        y: A target tensor. Shape :math:`(N, C, H, W)`.
+        y_parent: Flag to control dependency on previous layer of pyramid.
+        kernel_size: The side-length of the sliding window used in comparison for information content.
+        sigma_nsq: Parameter of visual distortion model.
+
+    Returns:
+        Information Content Maps.
+    """
+
+    EPS = torch.finfo(x.dtype).eps
+    n_channels = x.size(1)
+    kernel = average_filter2d(kernel_size=kernel_size).repeat(x.size(1), 1, 1, 1).to(x)
+    padding_up = kernel.size(-1) // 2
+    padding_down = kernel.size(-1) - padding_up
+
+    mu_x = F.conv2d(input=F.pad(x, pad=[padding_up, padding_down, padding_up, padding_down]), weight=kernel, padding=0,
+                    groups=n_channels)
+    mu_y = F.conv2d(input=F.pad(y, pad=[padding_up, padding_down, padding_up, padding_down]), weight=kernel, padding=0,
+                    groups=n_channels)
+
+    mu_xx = mu_x ** 2
+    mu_yy = mu_y ** 2
+    mu_xy = mu_x * mu_y
+
+    sigma_xx = F.conv2d(F.pad(x ** 2, pad=[padding_up, padding_down, padding_up, padding_down]), weight=kernel,
+                        stride=1, padding=0, groups=n_channels) - mu_xx
+    sigma_yy = F.conv2d(F.pad(y ** 2, pad=[padding_up, padding_down, padding_up, padding_down]), weight=kernel,
+                        stride=1, padding=0, groups=n_channels) - mu_yy
+    sigma_xy = F.conv2d(F.pad(x * y, pad=[padding_up, padding_down, padding_up, padding_down]), weight=kernel,
+                        stride=1, padding=0, groups=n_channels) - mu_xy
+
+    sigma_xx = F.relu(sigma_xx)
+    sigma_yy = F.relu(sigma_yy)
+
+    g = sigma_xy / (sigma_yy + EPS)
+    vv = sigma_xx - g * sigma_xy
+    g = g.masked_fill(sigma_yy < EPS, 0)
+    vv[sigma_yy < EPS] = sigma_xx[sigma_yy < EPS]
+    g = g.masked_fill(sigma_xx < EPS, 0)
+    vv = vv.masked_fill(sigma_xx < EPS, 0)
+
+    block = [kernel_size, kernel_size]
+
+    nblv = y.size(-2) - block[0] + 1
+    nblh = y.size(-1) - block[1] + 1
+    nexp = nblv * nblh
+    N = block[0] * block[1]
+
+    assert block[0] % 2 == 1 and block[1] % 2 == 1, f'Expected odd block dimensions, got {block}'
+
+    Ly = (block[0] - 1) // 2
+    Lx = (block[1] - 1) // 2
+
+    if y_parent is not None:
+        # upscale y_parent and cut to the size of y
+
+        y_parent_up = _image_enlarge(y_parent)[:, :, :y.size(-2), :y.size(-1)]
+        N = N + 1
+
+    Y = torch.zeros(y.size(0), y.size(1), nexp, N)
+
+    n = -1
+    for ny in range(-Ly, Ly + 1):
+        for nx in range(-Lx, Lx + 1):
+            n = n + 1
+            foo = _shift(y, [ny, nx])
+            foo = foo[:, :, Ly:Ly + nblv, Lx:Lx + nblh]
+            Y[..., n] = foo.flatten(start_dim=-2, end_dim=-1)
+
+    if y_parent is not None:
+        n = n + 1
+        foo = y_parent_up
+        foo = foo[:, :, Ly:Ly + nblv, Lx:Lx + nblh]
+        Y[..., n] = foo.flatten(start_dim=-2, end_dim=-1)
+
+    C_u = torch.matmul(Y.transpose(-2, -1), Y) / nexp
+
+    recommended_torch_version = _parse_version('1.7.0')
+    torch_version = _parse_version(torch.__version__)
+    if len(torch_version) != 0 and torch_version >= recommended_torch_version:
+        eig_values, eig_vectors = torch.linalg.eigh(C_u)
+    else:
+        eig_values, eig_vectors = torch.symeig(C_u, eigenvectors=True)
+
+    sum_eig_values = torch.sum(eig_values, dim=-1).view(y.size(0), y.size(1), 1, 1)
+    non_zero_eig_values_matrix = torch.diag_embed(eig_values * (eig_values > 0))
+    sum_non_zero_eig_values = torch.sum(non_zero_eig_values_matrix, dim=(-2, -1), keepdim=True)
+
+    L = non_zero_eig_values_matrix * sum_eig_values / (sum_non_zero_eig_values + (sum_non_zero_eig_values == 0))
+
+    C_u = torch.matmul(torch.matmul(eig_vectors, L), eig_vectors.transpose(-2, -1))
+
+    C_u_inv = torch.inverse(C_u)
+
+    ss = torch.matmul(Y, C_u_inv) * Y / N
+    ss = torch.sum(ss, dim=-1, keepdim=True)
+    ss = ss.view(y.size(0), y.size(1), nblv, nblh)
+    g = g[:, :, Ly: Ly + nblv, Lx: Lx + nblh]
+    vv = vv[:, :, Ly: Ly + nblv, Lx: Lx + nblh]
+
+    # Calculate mutual information
+    scaled_eig_values = torch.diagonal(L, offset=0, dim1=-2, dim2=-1).unsqueeze(2).unsqueeze(3)
+
+    iw_map = torch.sum(torch.log2(1 + ((vv.unsqueeze(-1) + (1 + g.unsqueeze(-1) * g.unsqueeze(-1)) * sigma_nsq)
+                                       * ss.unsqueeze(-1) * scaled_eig_values + sigma_nsq * vv.unsqueeze(-1)) / (
+                                              sigma_nsq * sigma_nsq)), dim=-1)
+
+    iw_map[iw_map < EPS] = 0
+
+    return iw_map
+
+
+def _image_enlarge(x: torch.Tensor) -> torch.Tensor:
+    r"""Custom bilinear upscaling of an image.
+    The function upscales an input image with upscaling factor 4x-3, adds padding on boundaries as difference
+    and downscaled by the factor of 2.
+
+    Args:
+        x: An input tensor. Shape :math:`(N, C, H, W)`.
+
+    Returns:
+        Upscaled tensor.
+    """
+    t1 = F.interpolate(x, size=(int(4 * x.size(-2) - 3), int(4 * x.size(-1) - 3)), mode='bilinear', align_corners=False)
+    t2 = torch.zeros([x.size(0), 1, 4 * x.size(-2) - 1, 4 * x.size(-1) - 1]).to(x)
+    t2[:, :, 1: -1, 1:-1] = t1
+    t2[:, :, 0, :] = 2 * t2[:, :, 1, :] - t2[:, :, 2, :]
+    t2[:, :, -1, :] = 2 * t2[:, :, -2, :] - t2[:, :, -3, :]
+    t2[:, :, :, 0] = 2 * t2[:, :, :, 1] - t2[:, :, :, 2]
+    t2[:, :, :, -1] = 2 * t2[:, :, :, -2] - t2[:, :, :, -3]
+    out = t2[:, :, ::2, ::2]
+    return out
+
+
+def _shift(x: torch.Tensor, shift: list) -> torch.Tensor:
+    r""" Circular shift 2D matrix samples by OFFSET (a [Y,X] 2-vector), such that  RES(POS) = MTX(POS-OFFSET).
+
+    Args:
+        x: An input tensor. Shape :math:`(N, C, H, W)`.
+        shift: Offset list.
+
+    Returns:
+        The circular shiftet tensor.
+    """
+    x_shifted = torch.cat((x[..., -shift[0]:, :], x[..., :-shift[0], :]), dim=-2)
+    x_shifted = torch.cat((x_shifted[..., -shift[1]:], x_shifted[..., :-shift[1]]), dim=-1)
+    return x_shifted
diff --git a/piq/ms_ssim.py b/piq/ms_ssim.py
index cf2431bf..fde41323 100644
--- a/piq/ms_ssim.py
+++ b/piq/ms_ssim.py
@@ -69,6 +69,8 @@ def multi_scale_ssim(x: torch.Tensor, y: torch.Tensor, kernel_size: int = 11, ke
     else:
         # Normalize scale weights
         scale_weights = (scale_weights / scale_weights.sum()).to(x)
+    if scale_weights.size(0) != scale_weights.numel():
+        raise ValueError(f'Expected a vector of weights, got {scale_weights.dim()}D tensor')
 
     kernel = gaussian_filter(kernel_size, kernel_sigma).repeat(x.size(1), 1, 1, 1).to(x)
 


Test patch:
diff --git a/tests/results_benchmark.py b/tests/results_benchmark.py
index dfe8f527..ad8e86fa 100644
--- a/tests/results_benchmark.py
+++ b/tests/results_benchmark.py
@@ -20,6 +20,7 @@
     "PSNR": functools.partial(piq.psnr, reduction='none'),
     "SSIM": functools.partial(piq.ssim, reduction='none'),
     "MS-SSIM": functools.partial(piq.multi_scale_ssim, reduction='none'),
+    "IW-SSIM": functools.partial(piq.information_weighted_ssim, reduction='none'),
     "VIFp": functools.partial(piq.vif_p, reduction='none'),
     "GMSD": functools.partial(piq.gmsd, reduction='none'),
     "MS-GMSD": functools.partial(piq.multi_scale_gmsd, reduction='none'),
@@ -168,7 +169,7 @@ def main(dataset_name: str, path: Path, metrics: List, batch_size: int, device:
 
     # Init dataset and dataloader
     dataset = DATASETS[dataset_name](root=path)
-    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=4)
+    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=0)
 
     for name in metrics:
         gt_scores, metric_scores = eval_metric(loader, METRICS[name], device=device)
diff --git a/tests/test_iw_ssim.py b/tests/test_iw_ssim.py
new file mode 100644
index 00000000..25f35b29
--- /dev/null
+++ b/tests/test_iw_ssim.py
@@ -0,0 +1,279 @@
+import torch
+import itertools
+import pytest
+from piq import InformationWeightedSSIMLoss, information_weighted_ssim
+from typing import Tuple, List
+from skimage.io import imread
+from contextlib import contextmanager
+
+
+@contextmanager
+def raise_nothing(enter_result=None):
+    yield enter_result
+
+
+@pytest.fixture(scope='module')
+def x_rand() -> torch.Tensor:
+    return torch.rand(3, 3, 161, 161)
+
+
+@pytest.fixture(scope='module')
+def y_rand() -> torch.Tensor:
+    return torch.rand(3, 3, 161, 161)
+
+
+@pytest.fixture(scope='module')
+def ones_zeros_4d() -> Tuple[torch.Tensor, torch.Tensor]:
+    return torch.ones(3, 3, 161, 161), torch.zeros(3, 3, 161, 161)
+
+
+@pytest.fixture(scope='module')
+def test_images() -> List[Tuple[torch.Tensor, torch.Tensor]]:
+    x_grey = torch.tensor(imread('tests/assets/goldhill_jpeg.gif')).unsqueeze(0).unsqueeze(0)
+    y_grey = torch.tensor(imread('tests/assets/goldhill.gif')).unsqueeze(0).unsqueeze(0)
+    x_rgb = torch.tensor(imread('tests/assets/I01.BMP')).permute(2, 0, 1).unsqueeze(0)
+    y_rgb = torch.tensor(imread('tests/assets/i01_01_5.bmp')).permute(2, 0, 1).unsqueeze(0)
+    return [(x_grey, y_grey), (x_rgb, y_rgb)]
+
+
+@pytest.fixture(scope='module')
+def scale_weights() -> torch.Tensor:
+    return torch.tensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
+
+
+# ================== Test function: `information_weighted_ssim` ==================
+def test_iw_ssim_measure_is_one_for_equal_tensors(x_rand: torch.Tensor, device: str) -> None:
+    x_rand = x_rand.to(device)
+    y_rand = x_rand.clone()
+    measure = information_weighted_ssim(x_rand, y_rand, data_range=1.)
+    assert torch.allclose(measure, torch.ones_like(measure)), \
+        f'If equal tensors are passed IW-SSIM must be equal to 1 ' \
+        f'(considering floating point operation error up to 1 * 10^-6), got {measure + 1}'
+
+
+def test_iw_ssim_reduction(x_rand: torch.Tensor, y_rand: torch.Tensor, device: str) -> None:
+    for mode in ['mean', 'sum', 'none']:
+        information_weighted_ssim(x_rand.to(device), y_rand.to(device), reduction=mode)
+
+    for mode in [None, 'n', 2]:
+        with pytest.raises(ValueError):
+            information_weighted_ssim(x_rand.to(device), y_rand.to(device), reduction=mode)
+
+
+def test_iw_ssim_raises_if_tensors_have_different_shapes(x_rand: torch.Tensor, y_rand: torch.Tensor,
+                                                         scale_weights: torch.Tensor, device: str) -> None:
+
+    dims = [[3], [2, 3], [160, 161], [160, 161]]
+
+    for size in list(itertools.product(*dims)):
+        wrong_shape_x = torch.rand(size).to(x_rand)
+        print(wrong_shape_x.size())
+        if wrong_shape_x.size() == x_rand.size():
+            information_weighted_ssim(wrong_shape_x.to(device), x_rand.to(device))
+        else:
+            with pytest.raises(AssertionError):
+                information_weighted_ssim(wrong_shape_x.to(device), x_rand.to(device))
+
+    information_weighted_ssim(x_rand.to(device), y_rand.to(device), scale_weights=scale_weights.to(device))
+
+    wrong_scale_weights = torch.rand(2, 2)
+    with pytest.raises(ValueError):
+        information_weighted_ssim(x_rand.to(device), y_rand.to(device), scale_weights=wrong_scale_weights.to(device))
+
+
+def test_iw_ssim_raises_if_tensors_have_different_types(x_rand: torch.Tensor, device: str) -> None:
+    wrong_type_x = list(range(10))
+    with pytest.raises(AssertionError):
+        information_weighted_ssim(wrong_type_x, x_rand.to(device))
+
+
+def test_iw_ssim_raises_if_kernel_size_greater_than_image(x_rand: torch.Tensor, y_rand: torch.Tensor,
+                                                          device: str) -> None:
+    kernel_size = 11
+    levels = 5
+    min_size = (kernel_size - 1) * 2 ** (levels - 1) + 1
+    wrong_size_x = x_rand[:, :, :min_size - 1, :min_size - 1]
+    wrong_size_y = y_rand[:, :, :min_size - 1, :min_size - 1]
+    with pytest.raises(ValueError):
+        information_weighted_ssim(wrong_size_x.to(device), wrong_size_y.to(device), kernel_size=kernel_size)
+
+
+@pytest.mark.parametrize(
+    "data_range", [128, 255],
+)
+def test_iw_ssim_supports_different_data_ranges(x_rand: torch.Tensor, y_rand: torch.Tensor, data_range: int,
+                                                device: str) -> None:
+
+    x_scaled = (x_rand * data_range).type(torch.uint8)
+    y_scaled = (y_rand * data_range).type(torch.uint8)
+
+    measure_scaled = information_weighted_ssim(x_scaled.to(device), y_scaled.to(device), data_range=data_range)
+    measure = information_weighted_ssim(
+        x_scaled.to(device) / float(data_range),
+        y_scaled.to(device) / float(data_range),
+        data_range=1.0
+    )
+    diff = torch.abs(measure_scaled - measure)
+    assert (diff <= 1e-6).all(), f'Result for same tensor with different data_range should be the same, got {diff}'
+
+
+def test_iw_ssim_fails_for_incorrect_data_range(x_rand: torch.Tensor, y_rand: torch.Tensor, device: str) -> None:
+    # Scale to [0, 255]
+    x_scaled = (x_rand * 255).type(torch.uint8)
+    y_scaled = (y_rand * 255).type(torch.uint8)
+    with pytest.raises(AssertionError):
+        information_weighted_ssim(x_scaled.to(device), y_scaled.to(device), data_range=1.0)
+
+
+@pytest.mark.parametrize(
+    "dtype", [torch.float32, torch.float64],
+)
+def test_iw_ssim_preserves_dtype(x_rand: torch.Tensor, y_rand: torch.Tensor, dtype: torch.dtype, device: str) -> None:
+    output = information_weighted_ssim(x_rand.to(device=device, dtype=dtype), y_rand.to(device=device, dtype=dtype))
+    assert output.dtype == dtype
+
+
+def test_iw_ssim_corresponds_to_matlab(test_images: List, device: str):
+    x_gray, y_gray = test_images[0]
+    x_rgb, y_rgb = test_images[1]
+    matlab_gray = torch.tensor(0.886297251092821, device=device)
+    matlab_rgb = torch.tensor(0.946804801436296, device=device)
+
+    score_gray = information_weighted_ssim(x_gray.to(device), y_gray.to(device), data_range=255)
+
+    assert torch.isclose(score_gray, matlab_gray, atol=1e-5),\
+        f'Expected {matlab_gray:.4f}, got {score_gray:.4f} for gray scale case.'
+
+    score_rgb = information_weighted_ssim(x_rgb.to(device), y_rgb.to(device), data_range=255)
+
+    assert torch.isclose(score_rgb, matlab_rgb, atol=1e-5),\
+        f'Expected {matlab_rgb:.8f}, got {score_rgb:.8f} for rgb case.'
+
+
+# ================== Test class: `InformationWeightedSSIMLoss` ==================
+def test_iw_ssim_loss_is_one_for_equal_tensors(x_rand: torch.Tensor, device: str) -> None:
+    x_rand = x_rand.to(device)
+    y_rand = x_rand.clone()
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    measure = loss(x_rand, y_rand)
+    assert torch.allclose(measure, torch.zeros_like(measure), atol=1e-5), \
+        f'If equal tensors are passed IW-SSIM must be equal to 0 ' \
+        f'(considering floating point operation error up to 1 * 10^-5), got {measure}'
+
+
+def test_iw_ssim_loss_reduction(x_rand: torch.Tensor, y_rand: torch.Tensor, device: str) -> None:
+    for mode in ['mean', 'sum', 'none']:
+        loss = InformationWeightedSSIMLoss(reduction=mode)
+        loss(x_rand.to(device), y_rand.to(device))
+
+    for mode in [None, 'n', 2]:
+        with pytest.raises(ValueError):
+            loss = InformationWeightedSSIMLoss(reduction=mode)
+            loss(x_rand.to(device), y_rand.to(device))
+
+
+def test_iw_ssim_loss_raises_if_tensors_have_different_shapes(x_rand: torch.Tensor, y_rand: torch.Tensor,
+                                                              scale_weights: torch.Tensor, device: str) -> None:
+
+    dims = [[3], [2, 3], [160, 161], [160, 161]]
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    for size in list(itertools.product(*dims)):
+        wrong_shape_x = torch.rand(size).to(x_rand)
+        print(wrong_shape_x.size())
+        if wrong_shape_x.size() == x_rand.size():
+            loss(wrong_shape_x.to(device), x_rand.to(device))
+        else:
+            with pytest.raises(AssertionError):
+                loss(wrong_shape_x.to(device), x_rand.to(device))
+
+    loss = InformationWeightedSSIMLoss(data_range=1., scale_weights=scale_weights)
+    loss(x_rand.to(device), y_rand.to(device))
+    wrong_scale_weights = torch.rand(2, 2)
+    loss = InformationWeightedSSIMLoss(data_range=1., scale_weights=wrong_scale_weights)
+    with pytest.raises(ValueError):
+        loss(x_rand.to(device), y_rand.to(device))
+
+
+def test_iw_ssim_loss_raises_if_tensors_have_different_types(x_rand: torch.Tensor, device: str) -> None:
+    wrong_type_x = list(range(10))
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    with pytest.raises(AssertionError):
+        loss(wrong_type_x, x_rand.to(device))
+
+
+def test_iw_ssim_loss_raises_if_kernel_size_greater_than_image(x_rand: torch.Tensor, y_rand: torch.Tensor,
+                                                               device: str) -> None:
+    kernel_size = 11
+    levels = 5
+    min_size = (kernel_size - 1) * 2 ** (levels - 1) + 1
+    wrong_size_x = x_rand[:, :, :min_size - 1, :min_size - 1]
+    wrong_size_y = y_rand[:, :, :min_size - 1, :min_size - 1]
+    loss = InformationWeightedSSIMLoss(data_range=1., kernel_size=kernel_size)
+    with pytest.raises(ValueError):
+        loss(wrong_size_x.to(device), wrong_size_y.to(device))
+
+
+@pytest.mark.parametrize(
+    "data_range", [128, 255],
+)
+def test_iw_ssim_loss_supports_different_data_ranges(x_rand: torch.Tensor, y_rand: torch.Tensor, data_range: int,
+                                                     device: str) -> None:
+
+    x_scaled = (x_rand * data_range).type(torch.uint8)
+    y_scaled = (y_rand * data_range).type(torch.uint8)
+
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    loss_scaled = InformationWeightedSSIMLoss(data_range=data_range)
+
+    measure_scaled = loss_scaled(x_scaled.to(device), y_scaled.to(device))
+    measure = loss(
+        x_scaled.to(device) / float(data_range),
+        y_scaled.to(device) / float(data_range)
+    )
+    diff = torch.abs(measure_scaled - measure)
+    assert (diff <= 1e-6).all(), f'Result for same tensor with different data_range should be the same, got {diff}'
+
+
+def test_iw_ssim_loss_fails_for_incorrect_data_range(x_rand: torch.Tensor, y_rand: torch.Tensor, device: str) -> None:
+    # Scale to [0, 255]
+    x_scaled = (x_rand * 255).type(torch.uint8)
+    y_scaled = (y_rand * 255).type(torch.uint8)
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    with pytest.raises(AssertionError):
+        loss(x_scaled.to(device), y_scaled.to(device))
+
+
+@pytest.mark.parametrize(
+    "dtype", [torch.float32, torch.float64],
+)
+def test_iw_ssim_loss_preserves_dtype(x_rand: torch.Tensor, y_rand: torch.Tensor, dtype: torch.dtype,
+                                      device: str) -> None:
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    output = loss(x_rand.to(device=device, dtype=dtype), y_rand.to(device=device, dtype=dtype))
+    assert output.dtype == dtype
+
+
+def test_iw_ssim_loss_corresponds_to_matlab(test_images: List, device: str):
+    x_gray, y_gray = test_images[0]
+    x_rgb, y_rgb = test_images[1]
+    matlab_gray = 1 - torch.tensor(0.886297251092821, device=device)
+    matlab_rgb = 1 - torch.tensor(0.946804801436296, device=device)
+
+    loss = InformationWeightedSSIMLoss(data_range=255)
+    score_gray = loss(x_gray.to(device), y_gray.to(device))
+
+    assert torch.isclose(score_gray, matlab_gray, atol=1e-5),\
+        f'Expected {matlab_gray:.8f}, got {score_gray:.8f} for gray scale case.'
+
+    score_rgb = loss(x_rgb.to(device), y_rgb.to(device))
+
+    assert torch.isclose(score_rgb, matlab_rgb, atol=1e-5),\
+        f'Expected {matlab_rgb:.8f}, got {score_rgb:.8f} for rgb case.'
+
+
+def test_iw_ssim_loss_backprop(x_rand: torch.Tensor, y_rand: torch.Tensor, device: str):
+    x_rand.requires_grad_(True)
+    loss = InformationWeightedSSIMLoss(data_range=1.)
+    score_gray = loss(x_rand.to(device), y_rand.to(device))
+    score_gray.backward()
+    assert torch.isfinite(x_rand.grad).all(), f'Expected finite gradient values, got {x_rand.grad}.'
diff --git a/tests/test_ms_ssim.py b/tests/test_ms_ssim.py
index f74083e4..33620b6f 100644
--- a/tests/test_ms_ssim.py
+++ b/tests/test_ms_ssim.py
@@ -76,6 +76,7 @@ def test_multi_scale_ssim_measure_is_less_or_equal_to_one(ones_zeros_4d_5d: Tupl
 
 
 def test_multi_scale_ssim_raises_if_tensors_have_different_shapes(x_y_4d_5d, device: str) -> None:
+    x = x_y_4d_5d[0].to(device)
     y = x_y_4d_5d[1].to(device)
     dims = [[3], [2, 3], [161, 162], [161, 162]]
     if y.dim() == 5:
@@ -88,7 +89,7 @@ def test_multi_scale_ssim_raises_if_tensors_have_different_shapes(x_y_4d_5d, dev
             with pytest.raises(AssertionError):
                 multi_scale_ssim(wrong_shape_x, y)
     scale_weights = torch.rand(2, 2)
-    with pytest.raises(AssertionError):
+    with pytest.raises(ValueError):
         multi_scale_ssim(x, y, scale_weights=scale_weights)
 
 
@@ -199,6 +200,7 @@ def test_multi_scale_ssim_loss_is_less_or_equal_to_one(ones_zeros_4d_5d: Tuple[t
 
 
 def test_multi_scale_ssim_loss_raises_if_tensors_have_different_shapes(x_y_4d_5d, device: str) -> None:
+    x = x_y_4d_5d[0].to(device)
     y = x_y_4d_5d[1].to(device)
     dims = [[3], [2, 3], [161, 162], [161, 162]]
     if y.dim() == 5:
@@ -212,7 +214,7 @@ def test_multi_scale_ssim_loss_raises_if_tensors_have_different_shapes(x_y_4d_5d
                 MultiScaleSSIMLoss()(wrong_shape_x, y)
 
     scale_weights = torch.rand(2, 2)
-    with pytest.raises(AssertionError):
+    with pytest.raises(ValueError):
         MultiScaleSSIMLoss(scale_weights=scale_weights)(x, y)