如何度量两幅图像的相似度--结构相似度SSIM 原理及代码

如何度量两幅图像的相似度--结构相似度SSIM

1. 什么是 SSIM

结构相似性指数(Structural Similarity Index measure,SSIM)用作度量两个给定图像之间的相似性。


image

如,这两个图像之间的 相似度=0.7816

SSIM 指标从图像中提取 3 个关键特征:

  • 亮度
  • 对比
  • 结构

两个图像之间的比较是基于这 3 个特征进行的


image

公式:
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计算原理参考链接

SSIM 有很多变体,如: Multi-Scale SSIM, Multi-component SSIM, Structural Dissimilarity,Complex Wavelet SSIM等等。维基百科可查看更多描述。

2. SSIM 有什么用

  1. 用做评价指标,度量两个图像的相似度。比如:做图像恢复/建模的时候,评价图像的恢复情况。
  2. 用做 loss function: 可用作分割以及图像重建的损失函数。不同的任务用法不一样。也是比较流行的方法。


    image

3. 使用 pytorch 计算 SSIM

这部分是本次的重点,会详细介绍 二维图像以及三维图像的SSIM计算方法, 以及如何用将其用做loss。

本部分实战简单,知识点丰富,内容有趣。一起来学习吧。

3.1 二维图像 SSIM 计算

2D SSIM pytorch code

这个代码是 pytorch 计算 SSIM 的标准方法,在github上搜索会发现,大家使用的都是这个方法。

这里我会带领大家一起实现官方的例子。

获得上述 GIF 结果大概需要如下技能:

  • SSIM 的计算
  • SSIM 用做 loss, 训练一个随机图像接近我们的目标图像(爱因斯坦肖像)
  • 将上述训练中,每个epoch生成的图像保存下来制作GIF
    • 用 Python 给图片加字
    • 用 Python 制作 GIF

任务很多,但是每个任务都很简单,我们开始吧~~


3.1.1 准备工作

下载 3.1 节中给出的地址中的代码,copy 整个 pytorch_ssim 文件夹。不知道如何单独下载一个文件夹,查看之前的教程。

Github 单文件快速下载 -- 使用 DownGit

这个文件夹其实就只有一个初始化文件。

image

如果还是不会下载,就自己新建一个同名文件夹(注意在 pycharm 中选择 新建 python package 而不是普通的文件夹),这样新建好后就会自动创建 __init__.py 文件. 然后把内容复制进去。
image

新建一个项目,把下载好的文件夹放进去。同时下载实验数据 ’einstein.png‘ 爱因斯坦肖像图。此时,你的项目里应该有打勾的3个文件。


image

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 保存图片?

这个链接提供了3种保存方法

到这里 官网的教程结束了,但是我们没有得到想要的 GIF!

只得到了每个 epoch 的图像


image

现在我们要做的就是:把每个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' 文件夹(如果没有,手动创建)

image

我们就得到了这样的一些图片。

3.2.2 制作GIF

我们把这些图片,按照epoch的顺序依次播放,并保存为GIF。

这里涉及到图片的排序问题。如果我们直接使用 sorted()函数排序

imgs = sorted(glob('SSIM_VALUE/*.png'))

会发现排出来的顺序跟我们要的顺序不一样。

image

而我们想要按数字大小排序,从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, 我们下篇博客见~

白天工作晚上写文,呕心沥血

觉得写的不错的话最后,求点赞,评论,收藏。或者一键三连

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 204,293评论 6 478
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,604评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,958评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,729评论 1 277
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,719评论 5 366
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,630评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,000评论 3 397
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,665评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,909评论 1 299
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,646评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,726评论 1 330
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,400评论 4 321
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,986评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,959评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,197评论 1 260
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 44,996评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,481评论 2 342

推荐阅读更多精彩内容