使用pytorch计算图像的Affinity Matrix

什么是Affinity Matrix

affinity matrix翻译过来是相关性矩阵，实际上就是计算图像中每个点与其他所有点的相关性

假设我们有一个三通道的图像，长为H，宽为W，那么每个像素点都是由3个分量构成的向量，我们可以计算任意两个点之间的向量的余弦值

$$cos<a,b>=\frac{a\ast b}{\mid a\mid \ast \mid b\mid }$$
当我们计算出来的余弦值为0的时候，表示两个向量垂直，这就表示他们之间相似度很低，当我们计算出来的余弦值为1的时候，就表示两个向量平行，他们之间非常相似
我们使用矩阵乘法来计算相似度矩阵，对于一幅大小为WH的图片，我们先将图片变成一维向量，然后将这个一维向量乘以它的转置，这样就得到长为WH，宽也为W*H的Affinity Matrix，这个矩阵是对称的。

详细代码如下

import torch
import matplotlib.pyplot as plt
import torchvision.transforms as tfs
from PIL import Image
import time


# img [3, width, height]
def getAffinity_Matrix(img):
    img = img.permute(1,2,0)
    # [width, height]
    affinity = torch.zeros(img.shape[0]*img.shape[1], img.shape[0]*img.shape[1])
    print(affinity.shape)
    img1 = img.reshape(-1, img.shape[-1])
    # 计算向量的模
    img_ = torch.sqrt((img1[:,:]**2).sum(dim=-1))
    img2 = img.reshape(-1, img.shape[-1])
    for idx in range(affinity.shape[1]):
        affinity[idx, :] = torch.mul(img1[idx, :], img2[:, :]).sum(dim=-1)
        affinity[idx, :] = affinity[idx, :]/img_[idx]
    for idx in range(affinity.shape[0]):
        #continue
        affinity[:, idx] = affinity[:, idx]/img_[idx]
    print(affinity)
    return affinity

def display(affinity):
    plt.imshow(affinity)
    plt.colorbar()
    plt.savefig("affinity.jpg")
    plt.show()

def process(img_root, rate=16):
    img = Image.open(img_root)
    size = 224//rate
    img = img.resize((size,size))
    plt.imshow(img)
    plt.savefig("tmp.jpg")
    img = tfs.ToTensor()(img)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    img = img.to(device)
    return img

if __name__ == "__main__":
    st = time.time()
    img = "a.jpg"
    img = process(img)
    affinity = getAffinity_Matrix(img)
    display(affinity)
    print("cost time: %.4fs"%(time.time()-st))

代码运行结果

我们输入一张小狗的图片得到如下结果

再输入一张图片看看结果