如何快速计算大量二进制图像的结构相似性指数（SSIM）

我的DataFrame列中存储了大量二进制图像，均为形状(224, 224)、数据类型np.uint8的numpy数组，示例为四边形的二进制图。

我通过df.sample随机选取一个图像作为输入样本，样本格式如下：

[[0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 ...
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 255 255 255]]

我的目标是给DataFrame新增一列，存储样本图像与所有其他图像的结构相似性指数（SSIM）。目前尝试了两个SSIM实现包，代码如下：

from image_similarity_measures.quality_metrics import ssim
from skimage.metrics import structural_similarity as ssim_2
import timeit
import numpy as np
import pandas as pd

sample= df['Images'].sample(n=1).values[0]

# Image Similarity Measures 包实现
start = timeit.default_timer()
df['Similarity_1'] = df['Images'].apply(lambda x: ssim(x, sample))
stop = timeit.default_timer()
print('Similarity Measures: ', stop - start)

# Skimage 包实现
start = timeit.default_timer()
df['Similarity_2'] = df['Images'].apply(lambda x: ssim_2(x, sample))
stop = timeit.default_timer()
print('Skimage: ', stop - start)

针对1000张图像的测试结果：

Similarity Measures: 21.4252301
Skimage: 3.3435048999999992

由于需要处理多达7万张图像，当前速度无法满足需求。之前用sklearn的余弦相似度速度很快，但准确性远低于SSIM。想请教：

1. 如何提升当前SSIM计算的代码速度？
2. 是否有更高效的二进制图像比较方法？

一、SSIM计算速度优化 #

1. 替换apply为向量化批量处理 #

pandas.apply本质是逐元素循环，效率极低。可以先将所有图像转为三维numpy数组(N, 224, 224)，再批量计算：

import numpy as np
from skimage.metrics import structural_similarity as ssim_2

# 转换所有图像为三维数组
all_images = np.stack(df['Images'].values)
sample = df['Images'].sample(n=1).values[0]

# 批量计算SSIM
def batch_ssim(images, ref):
    scores = []
    for img in images:
        scores.append(ssim_2(img, ref))
    return np.array(scores)

start = timeit.default_timer()
df['Similarity_optimized'] = batch_ssim(all_images, sample)
stop = timeit.default_timer()
print('Batch SSIM: ', stop - start)

2. 用Numba加速循环 #

针对二进制图像简化SSIM计算逻辑，再用Numba的JIT编译加速循环：

from numba import jit

@jit(nopython=True)
def ssim_binary(img1, img2):
    # 二进制图像仅含0和255，简化SSIM计算
    mu1 = img1.mean()
    mu2 = img2.mean()
    sigma1 = img1.var()
    sigma2 = img2.var()
    sigma12 = np.mean((img1 - mu1) * (img2 - mu2))
    C1 = (0.01 * 255)**2
    C2 = (0.03 * 255)**2
    numerator = (2 * mu1 * mu2 + C1) * (2 * sigma12 + C2)
    denominator = (mu1**2 + mu2**2 + C1) * (sigma1 + sigma2 + C2)
    return numerator / denominator

@jit(nopython=True)
def batch_ssim_numba(images, ref):
    n = images.shape[0]
    scores = np.zeros(n)
    for i in range(n):
        scores[i] = ssim_binary(images[i], ref)
    return scores

# 调用方式
all_images = np.stack(df['Images'].values)
df['Similarity_numba'] = batch_ssim_numba(all_images, sample)

3. 多进程并行计算 #

利用多核CPU拆分计算任务，适合超大样本量：

from concurrent.futures import ProcessPoolExecutor

def compute_ssim(img):
    return ssim_2(img, sample)

all_images = df['Images'].values
with ProcessPoolExecutor() as executor:
    scores = list(executor.map(compute_ssim, all_images))
df['Similarity_parallel'] = scores

二、二进制图像专属高效比较方法 #

1. IoU（交并比） #

针对二进制图像的前景区域匹配，计算速度极快：

# 批量计算IoU
all_images = np.stack(df['Images'].values)
sample_bool = sample == 255
intersections = np.logical_and(all_images == 255, sample_bool).sum(axis=(1,2))
unions = np.logical_or(all_images == 255, sample_bool).sum(axis=(1,2))
# 避免除以0
df['IoU'] = intersections / np.where(unions == 0, 1, unions)

2. 感知哈希+汉明距离 #

将图像转换为哈希值后计算距离，适合大规模检索场景：

def binary_hash(img):
    # 均值哈希：压缩图像后生成哈希
    img_resized = img[::8, ::8]  # 压缩为28x28
    hash_val = (img_resized > img_resized.mean()).flatten()
    return hash_val

# 预计算所有图像哈希
all_hashes = np.array([binary_hash(img) for img in df['Images'].values])
sample_hash = binary_hash(sample)

# 计算汉明距离并转为相似度
hamming_distances = np.sum(all_hashes != sample_hash, axis=1)
df['Hash_Similarity'] = 1 - (hamming_distances / len(sample_hash))

三、其他建议 #

• 提前将DataFrame中的图像列转为三维numpy数组，避免重复类型转换开销；
• 若需频繁检索，可预计算所有图像的特征（如简化SSIM统计量、哈希值）并存储，后续仅需对比样本特征即可。

内容的提问来源于stack exchange，提问作者Johnny