如何度量两幅图像的相似度--结构相似度 SSIM 原理及代码
Posted Tina姐
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了如何度量两幅图像的相似度--结构相似度 SSIM 原理及代码相关的知识，希望对你有一定的参考价值。
如何度量两幅图像的相似度–结构相似度SSIM
本文目录
文章目录
1. 什么是 SSIM
结构相似性指数（Structural Similarity Index measure，SSIM）用作度量两个给定图像之间的相似性。
如，这两个图像之间的 相似度=0.7816
SSIM 指标从图像中提取 3 个关键特征：
- 亮度
- 对比
- 结构
两个图像之间的比较是基于这 3 个特征进行的
公式：
S
S
I
M
(
x
,
y
)
=
(
2
μ
x
μ
y
+
c
1
)
(
2
σ
x
y
+
c
2
)
(
μ
x
2
)
+
μ
y
2
+
c
1
)
(
σ
x
2
+
σ
y
2
+
c
2
)
SSIM(x, y) = \\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)}
SSIM(x,y)=(μx2)+μy2+c1)(σx2+σy2+c2)(2μxμy+c1)(2σxy+c2)
公式详情见 维基百科
SSIM 有很多变体，如： Multi-Scale SSIM， Multi-component SSIM， Structural Dissimilarity，Complex Wavelet SSIM等等。维基百科可查看更多描述。
2. SSIM 有什么用
- 用做评价指标，度量两个图像的相似度。比如：做图像恢复/建模的时候，评价图像的恢复情况。
- 用做 loss function: 可用作分割以及图像重建的损失函数。不同的任务用法不一样。也是比较流行的方法。
3. 使用 pytorch 计算 SSIM
这部分是本次的重点，会详细介绍 二维图像以及三维图像的SSIM计算方法， 以及如何用将其用做loss。
本部分实战简单，知识点丰富，内容有趣。一起来学习吧。
3.1 二维图像 SSIM 计算
这个代码是 pytorch 计算 SSIM 的标准方法，在github上搜索会发现，大家使用的都是这个方法。
这里我会带领大家一起实现官方的例子。
获得上述 GIF 结果大概需要如下技能：
- SSIM 的计算
- SSIM 用做 loss, 训练一个随机图像接近我们的目标图像(爱因斯坦肖像)
- 将上述训练中，每个epoch生成的图像保存下来制作GIF
- 用 Python 给图片加字
- 用 Python 制作 GIF
任务很多，但是每个任务都很简单，我们开始吧~~
3.1.1 准备工作
下载 3.1 节中给出的地址中的代码，copy 整个 pytorch_ssim 文件夹。不知道如何单独下载一个文件夹，查看之前的教程。
这个文件夹其实就只有一个初始化文件。
如果还是不会下载，就自己新建一个同名文件夹（注意在 pycharm 中选择 新建 python package 而不是普通的文件夹），这样新建好后就会自动创建 __init__.py 文件. 然后把内容复制进去。
新建一个项目，把下载好的文件夹放进去。同时下载实验数据 ’einstein.png‘ 爱因斯坦肖像图。此时，你的项目里应该有打勾的3个文件。
3.1.2 官网的第一个案例
第一个案例：随机生成两个二维图像，并计算他们的 SSIM。
import pytorch_ssim
import torch
from torch.autograd import Variable
img1 = Variable(torch.rand(1, 1, 256, 256))
img2 = Variable(torch.rand(1, 1, 256, 256))
if torch.cuda.is_available():
img1 = img1.cuda()
img2 = img2.cuda()
print(pytorch_ssim.ssim(img1, img2))
ssim_loss = pytorch_ssim.SSIM(window_size = 11)
print(ssim_loss(img1, img2))
这里
import pytorch_ssim就是我们copy下来的文件夹 调用
pytorch_ssim.ssim直接计算二者的相似度 调用
pytorch_ssim.SSIM大写的SSIM是计算loss,但是二者的计算方法是一样的，只是写法不一样。
3.1.3 官网的第二个案例
第二个案例： 使用 SSIM 作为loss, 把随机数训练成目标图像（爱因斯坦）。
这里和官网的代码有些略微差别，官网的copy下来是有BUG的。改过的行 我会在后面 加 ###
import pytorch_ssim
import torch
from torch.autograd import Variable
from torch import optim
import cv2
import numpy as np
npImg1 = cv2.imread("einstein.png")
img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0) / 255.0
img2 = torch.rand(img1.size())
if torch.cuda.is_available():
img1 = img1.cuda()
img2 = img2.cuda()
img1 = Variable(img1, requires_grad=False)
img2 = Variable(img2, requires_grad=True)
# Functional: pytorch_ssim.ssim(img1, img2, window_size = 11, size_average = True)
ssim_value = pytorch_ssim.ssim(img1, img2).item() ###
print("Initial ssim:", ssim_value)
# Module: pytorch_ssim.SSIM(window_size = 11, size_average = True)
ssim_loss = pytorch_ssim.SSIM()
optimizer = optim.Adam([img2], lr=0.01)
epoch = 1 ###
while ssim_value < 0.95:
optimizer.zero_grad()
ssim_out = -ssim_loss(img1, img2)
ssim_value = round(- ssim_out.item(), 4) ### round 保留4位小数
# save image
vutils.save_image(img2, f'SSIMPNG/einstein_{epoch}_{ssim_value}.png') ### 官网教程没有保存每一次迭代的图像，这里我们保存下来。 如果没有SSIMPNG这个文件夹，手动创建一个。
print(ssim_value)
ssim_out.backward()
optimizer.step()
epoch += 1 ###
这部分代码也比较简单，就是当 SSIM 值小于 0.95 时一直循环迭代。中途保存每次迭代的图像，之后我们可以查看图像的变化过程。
这里另外涉及到一个知识点： 如何使用 pytorch 保存图片？
到这里 官网的教程结束了，但是我们没有得到想要的 GIF！
只得到了每个 epoch 的图像
现在我们要做的就是：把每个epoch的SSIM 值写在图像上，然后把这些图像做成GIF。
3.2 在图片上写字，并制作GIF
3.2.1 使用Python在图片上写字
主要用到的库： pillow
from PIL import Image, ImageDraw, ImageFont
# pip install pillow
def drawssim(filepath, value):
# 打开图像
im = Image.open(filepath)
# 告诉系统，你要在图像上画画了
draw = ImageDraw.Draw(im)
# 设置字体，大小=50
font = ImageFont.truetype('/usr/share/fonts/truetype/Sarai/Sarai.ttf', 40)
# 如何找系统上自带的字体： 百度搜索Linux字体存放位置
# 写内容, 初始位置(30, 10) 颜色红色
draw.text((30, 10), f'ssim = {value}', fill='red', font=font)
# 保存
im.save(filepath.replace('SSIMPNG', 'SSIM_VALUE'))
if __name__ == '__main__':
from glob import glob
imgs = glob('SSIMPNG/*.png')
for img in imgs:
filepath = img
value = img.split('_')[-1].split('.png')[0]
drawssim(filepath, value)
drawssim函数接收一个单独的文件地址和SSIM值，把值写在该图像上，并且保存到 ‘SSIM_VALUE’ 文件夹（如果没有，手动创建）
我们就得到了这样的一些图片。
3.2.2 制作GIF
我们把这些图片，按照epoch的顺序依次播放，并保存为GIF。
这里涉及到图片的排序问题。如果我们直接使用 sorted()函数排序
imgs = sorted(glob('SSIM_VALUE/*.png'))
会发现排出来的顺序跟我们要的顺序不一样。
而我们想要按数字大小排序，从1–76.
因此，我们需要自写一个排序算法。
import imageio # pip install imageio
# create gif
def create_gif(image_list, gif_name, duration=1.5):
'''
image_list: 图片列表，文件名
gif_name： 生成的 GIF 文件名
duratuion： 图像间隔时间
'''
frames = []
for img in image_list:
frames.append(imageio.imread(img))
imageio.mimsave(gif_name, frames, 'GIF', duration=duration)
class sortimg():
def tryint(self, s):
try:
return int(s)
except ValueError:
return s
def str2int(self, v_str):
return [self.tryint(sub_str) for sub_str in re.split('([0-9]+)', v_str)]
def sort_humanly(self, v_list):
return sorted(v_list, key=self.str2int)
if __name__ == '__main__':
from glob import glob
imgs = glob('SSIM_VALUE/*.png')
# 排序
sorted_imgs = sortimg().sort_humanly(imgs)
create_gif(sorted_imgs, 'ssim.gif', duration=0.1)
sortimg()对列表进行排序(按数字大小)， create_gif()使用imageio制作GIF。
如此我们就得到了之前那个 GIF 图像。官网教程实验已经完美复现啦。
3.3 3D 图像的 SSIM 计算和 loss
上述代码只提供了二维图像的计算方法，那如何计算三维呢？
这里，copy 另一个代码。
3D 代码地址
import pytorch_ssim
import torch
from torch.autograd import Variable
img1 = Variable(torch.rand(1, 1, 256, 256, 256))
img2 = Variable(torch.rand(1, 1, 256, 256, 256))
if torch.cuda.is_available():
img1 = img1.cuda()
img2 = img2.cuda()
print(pytorch_ssim.ssim3D(img1, img2))
ssim_loss = pytorch_ssim.SSIM3D(window_size = 11)
print(ssim_loss(img1, img2))
这个代码在原来的基础上新增了几个 3D 函数，同 2D 使用方法一样。
文章持续更新，可以关注微信公众号【医学图像人工智能实战营】获取最新动态，一个关注于医学图像处理领域前沿科技的公众号。坚持已实践为主，手把手带你做项目，打比赛，写论文。凡原创文章皆提供理论讲解，实验代码，实验数据。只有实践才能成长的更快，关注我们，一起学习进步~
我是Tina, 我们下篇博客见~
白天工作晚上写文，呕心沥血
觉得写的不错的话最后，求点赞，评论，收藏。或者一键三连
以上是关于如何度量两幅图像的相似度--结构相似度 SSIM 原理及代码的主要内容，如果未能解决你的问题，请参考以下文章
基于欧氏距离（或任何其他距离计算技术）提取的SIFT描述符估计两幅图像的相似度得分