mssim.pytorch

SSIM:

$$\begin{align} l(\mathbf{x}, \mathbf{y}) & = \frac{2\mu_x\mu_y+C_1}{\mu_x^2+\mu_y^2+C_1}, \quad C_1=(K_1L)^2, K_1=0.01, \\ c(\mathbf{x}, \mathbf{y}) & = \frac{2\sigma_{x}\sigma_{y}+C_2}{\sigma_x^2+\sigma_y^2+C_2}, \quad C_2=(K_2L)^2, K_2=0.03, \\ s(\mathbf{x}, \mathbf{y}) & = \frac{\sigma_{xy}+C_3}{\sigma_x\sigma_y+C_3}, \quad C_3=C_2/2, \\ \text{SSIM}(\mathbf{x}, \mathbf{y}) & = [l(\mathbf{x}, \mathbf{y})]^\alpha \cdot [c(\mathbf{x}, \mathbf{y})]^\beta \cdot [s(\mathbf{x}, \mathbf{y})]^\gamma \\ & = \frac{(2\mu_x\mu_y+C_1)(2\sigma_{xy}+C_2)}{(\mu_x^2+\mu_y^2+C_1)(\sigma_x^2+\sigma_y^2+C_2)} \\ & \alpha=\beta=\gamma=1. \end{align}$$

Multi Scale SSIM:

$$\begin{align} \text{MS-SSIM}(\mathbf{x}, \mathbf{y}) & = [l(\mathbf{x}, \mathbf{y})]^{\alpha_{M}} \cdot \prod^{M}_{j=1} [c_j(\mathbf{x}, \mathbf{y})]^{\beta_j} \cdot [s_j(\mathbf{x}, \mathbf{y})]^{\gamma_j}, \\ (M=5) \Rightarrow & \beta_1=\gamma_1=0.0448, \\ & \beta_2=\gamma_2=0.2856, \\ & \beta_3=\gamma_3=0.3001, \\ & \beta_4=\gamma_4=0.2363, \\ & \alpha_5=\beta_5=\gamma_5=0.1333. \end{align}$$

Weighted/ROI SSIM:

$$\begin{align} \text{WeightedSSIM}_{\mathbf{w}}(\mathbf{x}, \mathbf{y}) & = \frac{\sum_i w_i\text{SSIM}_i(\mathbf{x}, \mathbf{y})}{\sum_i w_i}. \end{align}$$

Gradient-based SSIM:

$$\begin{align} \nabla_h\mathbf{x} & = \sqrt{(\mathbf{x}*h_x)^2+(\mathbf{x}*h_y)^2}, \\ \text{GradientSSIM}(\mathbf{x}, \mathbf{y}) & = \text{SSIM}(\nabla_h\mathbf{x}, \nabla_h\mathbf{y}). \end{align}$$

Three-Component Structural Similarity (3-SSIM):

$$\begin{align} \text{ThreeSSIM}(\mathbf{x}, \mathbf{y}) & = \sum_{r\in \{\text{edge}, \text{texture}, \text{smooth}\} }\lambda_r\frac{\sum_i m_{ri}\text{SSIM}_i(\mathbf{x}, \mathbf{y})}{\sum_i m_{ri}}, \\ & \lambda_\text{edge}=0.5, \quad \lambda_\text{texture}=0.25, \quad \lambda_\text{smooth}=0.25. \end{align}$$

A better pytorch-based implementation for the mean structural similarity (MSSIM).

Compared to this widely used implementation: https://github.com/Po-Hsun-Su/pytorch-ssim, I further optimized and refactored the code.

At the same time, in this implementation, I have dealt with the problem that the calculation with the fp16 mode cannot be consistent with the calculation with the fp32 mode. Typecasting is used here to ensure that the computation is done in fp32 mode. This might also avoid unexpected results when using it as a loss.

Setting	SSIM1d	SSIM2d	SSIM3d	MS-SSIM2d	MS-SSIM3d (only pooling in the spatial domain)	WeightedSSIM	GradientSSIM	ThreeSSIM
data_dim	1	2 (Default)	3	2	3	1, 2 (Default), 3	2	2
forward return_msssim	False	False	False	True	True	False	False	False
window_size	int, [int]	int, [int, int]	int, [int, int, int]	int, [int, int]	int, [int, int, int]	int, [int] by data_dim	int, [int, int]	int, [int, int]
padding	int, [int]	int, [int, int]	int, [int, int, int]	int, [int, int]	int, [int, int, int]	int, [int] by data_dim	int, [int, int]	int, [int, int]
sigma	float, [float]	float, [float, float]	float, [float, float, float]	float, [float, float]	float, [float, float, float]	float, [float] by data_dim	float, [float, float]	float, [float, float]
in_channels	int	int	int	int	int	int	int	int
L	1, 255	1, 255	1, 255	1, 255	1, 255	1, 255	1, 255	1, 255
reduction	✅	✅	✅	✅	✅	✅	✅	✅
return_log	✅	✅	✅	❌	❌	✅	❌	❌
ensemble_kernel	✅	✅	✅	✅	✅	✅	✅	✅
weights	optional	optional	optional	❌	❌	required	❌	auto region masks
eps	optional for weights	optional for weights	optional for weights	❌	❌	✅	✅	✅

Current API

The SSIM family now uses a composable internal pipeline: preprocess -> SSIMCore -> reducer.

• Standard SSIM uses IdentityPreprocess and SSIMMapReducer;
• GradientSSIM swaps in GradientSSIMPreprocess;
• ThreeSSIM swaps in ThreeSSIMReducer.

This keeps the local SSIM map computation shared and makes new variants additive instead of duplicating the whole metric.

from ssim import (
    SSIM,
    MSSSIM,
    WeightedSSIM,
    GradientSSIM,
    ThreeSSIM,
    ssim,
    ms_ssim,
    weighted_ssim,
    gradient_ssim,
    three_ssim,
)

Group	Functional API	Module API
SSIM / Multi-Scale SSIM	ssim, ms_ssim	SSIM, MSSSIM
Weighted/ROI SSIM	ssim(..., weights=...), weighted_ssim	SSIM(...)(x, y, weights=...), WeightedSSIM
Gradient-based SSIM	gradient_ssim	GradientSSIM
Three-Component Structural Similarity (3-SSIM)	three_ssim	ThreeSSIM

score = SSIM(reduction="batch")(x, y)
ms_score = SSIM(reduction="batch")(x, y, return_msssim=True)
weighted_score = weighted_ssim(x, y, weights, window_size=5, L=1, reduction="batch")
gradient_score = GradientSSIM(window_size=5, L=1, reduction="batch")(x, y)
three_score = ThreeSSIM(window_size=5, L=1, reduction="batch", variance_window_size=5)(x, y)

Structural similarity index

When comparing images, the mean squared error (MSE)–while simple to implement–is not highly indicative of perceived similarity. Structural similarity aims to address this shortcoming by taking texture into account. More details can be seen at https://scikit-image.org/docs/dev/auto_examples/transform/plot_ssim.html?highlight=structure+similarity

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from lartpang_ssim import SSIM
from po_hsun_su_ssim import SSIM as PoHsunSuSSIM
from vainf_ssim import MS_SSIM as VainFMSSSIM
from vainf_ssim import SSIM as VainFSSIM
from skimage import data, img_as_float

img = img_as_float(data.camera())
rows, cols = img.shape

noise = np.ones_like(img) * 0.3 * (img.max() - img.min())
rng = np.random.default_rng()
noise[rng.random(size=noise.shape) > 0.5] *= -1

img_noise = img + noise
img_const = np.zeros_like(img)

img_tensor = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float()
img_noise_tensor = torch.from_numpy(img_noise).unsqueeze(0).unsqueeze(0).float()
img_const_tensor = torch.from_numpy(img_const).unsqueeze(0).unsqueeze(0).float()

fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 7))
ax = axes.ravel()

mse_none = F.mse_loss(img_tensor, img_tensor, reduction="mean")
mse_noise = F.mse_loss(img_tensor, img_noise_tensor, reduction="mean")
mse_const = F.mse_loss(img_tensor, img_const_tensor, reduction="mean")

# https://github.com/VainF/pytorch-msssim
vainf_ssim_none = VainFSSIM(channel=1, data_range=1)(img_tensor, img_tensor)
vainf_ssim_noise = VainFSSIM(channel=1, data_range=1)(img_tensor, img_noise_tensor)
vainf_ssim_const = VainFSSIM(channel=1, data_range=1)(img_tensor, img_const_tensor)
vainf_ms_ssim_none = VainFMSSSIM(channel=1, data_range=1)(img_tensor, img_tensor)
vainf_ms_ssim_noise = VainFMSSSIM(channel=1, data_range=1)(img_tensor, img_noise_tensor)
vainf_ms_ssim_const = VainFMSSSIM(channel=1, data_range=1)(img_tensor, img_const_tensor)

# use the settings of https://github.com/VainF/pytorch-msssim
ssim_vainf = SSIM(L=1, padding=0, ensemble_kernel=False)
ssim_none_0 = ssim_vainf(img_tensor, img_tensor)
ssim_noise_0 = ssim_vainf(img_tensor, img_noise_tensor)
ssim_const_0 = ssim_vainf(img_tensor, img_const_tensor)
ms_ssim_none_0 = ssim_vainf(img_tensor, img_tensor, return_msssim=True)
ms_ssim_noise_0 = ssim_vainf(img_tensor, img_noise_tensor, return_msssim=True)
ms_ssim_const_0 = ssim_vainf(img_tensor, img_const_tensor, return_msssim=True)

# https://github.com/Po-Hsun-Su/pytorch-ssim
pohsunsu_ssim_none = PoHsunSuSSIM()(img_tensor, img_tensor)
pohsunsu_ssim_noise = PoHsunSuSSIM()(img_tensor, img_noise_tensor)
pohsunsu_ssim_const = PoHsunSuSSIM()(img_tensor, img_const_tensor)

# use the settings of https://github.com/Po-Hsun-Su/pytorch-ssim
ssim_pohsunsu = SSIM(L=1, padding=None, ensemble_kernel=True)
ssim_none_1 = ssim_pohsunsu(img_tensor, img_tensor)
ssim_noise_1 = ssim_pohsunsu(img_tensor, img_noise_tensor)
ssim_const_1 = ssim_pohsunsu(img_tensor, img_const_tensor)


ax[0].imshow(img, cmap=plt.cm.gray, vmin=0, vmax=1)
ax[0].set_xlabel(
    f"MSE: {mse_none:.6f}\n"
    f"SSIM {ssim_none_0:.6f}, MS-SSIM {ms_ssim_none_0:.6f}\n"
    f"(VainF) SSIM: {vainf_ssim_none:.6f}, MS-SSIM {vainf_ms_ssim_none:.6f}\n"
    f"SSIM {ssim_none_1:.6f}\n"
    f"(PoHsunSu) SSIM: {pohsunsu_ssim_none:.6f}\n"
)
ax[0].set_title("Original image")

ax[1].imshow(img_noise, cmap=plt.cm.gray, vmin=0, vmax=1)
ax[1].set_xlabel(
    f"MSE: {mse_noise:.6f}\n"
    f"SSIM {ssim_noise_0:.6f}, MS-SSIM {ms_ssim_noise_0:.6f}\n"
    f"(VainF) SSIM: {vainf_ssim_noise:.6f}, MS-SSIM {vainf_ms_ssim_noise:.6f}\n"
    f"SSIM {ssim_noise_1:.6f}\n"
    f"(PoHsunSu) SSIM: {pohsunsu_ssim_noise:.6f}\n"
)
ax[1].set_title("Image with noise")

ax[2].imshow(img_const, cmap=plt.cm.gray, vmin=0, vmax=1)
ax[2].set_xlabel(
    f"MSE: {mse_const:.6f}\n"
    f"SSIM {ssim_const_0:.6f}, MS-SSIM {ms_ssim_const_0:.6f}\n"
    f"(VainF) SSIM: {vainf_ssim_const:.6f}, MS-SSIM {vainf_ms_ssim_const:.6f}\n"
    f"SSIM {ssim_const_1:.6f}\n"
    f"(PoHsunSu) SSIM: {pohsunsu_ssim_const:.6f}\n"
)
ax[2].set_title("Image plus constant")


[ax[i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[]) for i in range(len(axes))]

plt.tight_layout()
plt.savefig("results.png")

More Examples

# setting 4: for 4d float tensors with the data range [0, 1] and 1 channel, return the logarithmic form per batch item
ssim_caller = SSIM(return_log=True, reduction="batch").cuda()

# two 4d tensors
x = torch.randn(3, 1, 100, 100).cuda()
y = torch.randn(3, 1, 100, 100).cuda()
ssim_score_0 = ssim_caller(x, y)
# or in the fp16 mode (we have fixed the computation progress into the float32 mode to avoid the unexpected result)
with torch.cuda.amp.autocast(enabled=True):
    ssim_score_1 = ssim_caller(x, y)
assert torch.allclose(ssim_score_0, ssim_score_1)
print(ssim_score_0.shape, ssim_score_1.shape)

As A Loss

As you can see from the respective thresholds of the two cases below, it is easier to optimize towards MSSIM=1 than MSSIM=-1.

Optimize towards MSSIM=1

import matplotlib.pyplot as plt
import torch
from pytorch_ssim import SSIM
from skimage import data
from torch import optim

original_image = data.moon() / 255
target_image = torch.from_numpy(original_image).unsqueeze(0).unsqueeze(0).float().cuda()
predicted_image = torch.zeros_like(
    target_image, device=target_image.device, dtype=target_image.dtype, requires_grad=True
)
initial_image = predicted_image.clone()

ssim = SSIM().cuda()
initial_ssim_value = ssim(predicted_image, target_image)

ssim_value = initial_ssim_value
optimizer = optim.Adam([predicted_image], lr=0.01)
loss_curves = []
while ssim_value < 0.999:
    ssim_out = 1 - ssim(predicted_image, target_image)
    loss_curves.append(ssim_out.item())
    ssim_value = 1 - ssim_out.item()
    print(ssim_value)
    ssim_out.backward()
    optimizer.step()
    optimizer.zero_grad()

fig, axes = plt.subplots(nrows=2, ncols=4, figsize=(8, 4))
ax = axes.ravel()

ax[0].imshow(original_image, cmap=plt.cm.gray, vmin=0, vmax=1)
ax[0].set_title("Original Image")

ax[1].imshow(initial_image.squeeze().detach().cpu().numpy(), cmap=plt.cm.gray, vmin=0, vmax=1)
ax[1].set_xlabel(f"SSIM: {initial_ssim_value:.5f}")
ax[1].set_title("Initial Image")

ax[2].imshow(predicted_image.squeeze().detach().cpu().numpy(), cmap=plt.cm.gray, vmin=0, vmax=1)
ax[2].set_xlabel(f"SSIM: {ssim_value:.5f}")
ax[2].set_title("Predicted Image")

ax[3].plot(loss_curves)
ax[3].set_title("SSIM Loss Curve")

ax[4].set_title("Original Image")
ax[4].hist(original_image.ravel(), bins=256)
ax[4].ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))
ax[4].set_xlabel("Pixel Intensity")

ax[5].set_title("Initial Image")
ax[5].hist(initial_image.squeeze().detach().cpu().numpy().ravel(), bins=256)
ax[5].ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))
ax[5].set_xlabel("Pixel Intensity")

ax[6].set_title("Predicted Image")
ax[6].hist(predicted_image.squeeze().detach().cpu().numpy().ravel(), bins=256)
ax[6].ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))
ax[6].set_xlabel("Pixel Intensity")

plt.tight_layout()
plt.savefig("prediction.png")

Optimize towards MSSIM=-1

import matplotlib.pyplot as plt
import torch
from pytorch_ssim import SSIM
from skimage import data
from torch import optim

original_image = data.moon() / 255
target_image = torch.from_numpy(original_image).unsqueeze(0).unsqueeze(0).float().cuda()
predicted_image = torch.zeros_like(
    target_image, device=target_image.device, dtype=target_image.dtype, requires_grad=True
)
initial_image = predicted_image.clone()

ssim = SSIM(L=original_image.max() - original_image.min()).cuda()
initial_ssim_value = ssim(predicted_image, target_image)

ssim_value = initial_ssim_value
optimizer = optim.Adam([predicted_image], lr=0.01)
loss_curves = []
while ssim_value > -0.94:
    ssim_out = ssim(predicted_image, target_image)
    loss_curves.append(ssim_out.item())
    ssim_value = ssim_out.item()
    print(ssim_value)
    ssim_out.backward()
    optimizer.step()
    optimizer.zero_grad()

fig, axes = plt.subplots(nrows=2, ncols=4, figsize=(8, 4))
ax = axes.ravel()

ax[0].imshow(original_image, cmap=plt.cm.gray, vmin=0, vmax=1)
ax[0].set_title("Original Image")

ax[1].imshow(initial_image.squeeze().detach().cpu().numpy(), cmap=plt.cm.gray, vmin=0, vmax=1)
ax[1].set_xlabel(f"SSIM: {initial_ssim_value:.5f}")
ax[1].set_title("Initial Image")

ax[2].imshow(predicted_image.squeeze().detach().cpu().numpy(), cmap=plt.cm.gray, vmin=0, vmax=1)
ax[2].set_xlabel(f"SSIM: {ssim_value:.5f}")
ax[2].set_title("Predicted Image")

ax[3].plot(loss_curves)
ax[3].set_title("SSIM Loss Curve")

ax[4].set_title("Original Image")
ax[4].hist(original_image.ravel(), bins=256)
ax[4].ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))
ax[4].set_xlabel("Pixel Intensity")

ax[5].set_title("Initial Image")
ax[5].hist(initial_image.squeeze().detach().cpu().numpy().ravel(), bins=256)
ax[5].ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))
ax[5].set_xlabel("Pixel Intensity")

ax[6].set_title("Predicted Image")
ax[6].hist(predicted_image.squeeze().detach().cpu().numpy().ravel(), bins=256)
ax[6].ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))
ax[6].set_xlabel("Pixel Intensity")

plt.tight_layout()
plt.savefig("prediction.png")

Reference

• https://github.com/Po-Hsun-Su/pytorch-ssim
• https://github.com/VainF/pytorch-msssim
• https://scikit-image.org/docs/dev/auto_examples/transform/plot_ssim.html?highlight=structure+similarity
• Z. Wang, A. C. Bovik, H. R. Sheikh and E. P. Simoncelli, “Image quality assessment: From error visibility to structural similarity,” IEEE Transactions on Image Processing, vol. 13, no. 4, pp. 600-612, Apr. 2004.

Cite

If you find this library useful, please cite our bibtex:

@online{mssim.pytorch,
    author="lartpang",
    title="{A better pytorch-based implementation for the mean structural similarity. Differentiable simpler SSIM and MS-SSIM.}",
    url="https://github.com/lartpang/mssim.pytorch",
    note="(Jun 21, 2022)",
}