【以图搜图】Python实现根据图片批量匹配(查找)相似图片

祁彧w

已于 2023-08-01 19:43:44 修改

阅读量6.2k

点赞数 5

分类专栏： python 文章标签：算法图像检索以图搜图 python

于 2023-08-01 19:41:59 首次发布

本文链接：https://blog.csdn.net/Joker_Q/article/details/132048333

版权

python 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

文章介绍了一种方法，通过计算MSE和SSIM来比较查询图像与数据库图像的相似度，实现在本地根据图片查找相似图片的功能。在32x32图像中，平均检索速度约为34.64秒，且支持设置权重阈值提前结束。代码使用了Python的OpenCV、PIL、skimage和PyTorch库。

摘要由CSDN通过智能技术生成

目的：可以解决在本地实现根据图片查找相似图片的功能

背景：由于需要查找别人代码保存的图像的命名，但由于数据集是cifa10图像又小又多，所以直接找很费眼睛，所以实现用该代码根据图像查找图像，从而得到保存图像的命名。

方法：

1、将需要查找的图像(查询图像, queryImg)放入queryImgs文件夹，以及一个存放数据库图像的文件夹datasetImgs

2、批量读取查询图像

3、根据MSE(均方误差)和SSIM(结构相似性指数)计算权重，来比较两张图像的相似程度。

其中：

MSE(均方误差)：计算两张图片的每个像素值之间的平均差值，结果越小表示两张图片越相似。

SSIM(结构相似性指数)：比较两张图片的结构、亮度和对比度等方面的相似程度，结果介于-1到1之间，越接近1表示两张图片越相似。

$weight=MSE\times (1-SSIM)$

4、以224×224的大小显示当前queryImg和bestImg(数据库图像中相似度最高的图像)，title为queryImg和bestImg的文件名。

5、将bestImg移动到命名为“dstImgs”的文件夹，并保留源文件名称。

运行速度：

在1万张32×32图像中，平均检索速度：34.64s左右（不包括对数据库图像使用transform统一大小），可以为权重设置阈值提前结束检索（建议阈值为小于10）。

限制：

1、需要明确查询图像和数据库图像的大小，并手动更改对应注释的代码。

2、检索结果唯一，不能检索到多个结果（由于我知道我的datasets里只有唯一对应的图像，所以代码逻辑是只保存最相似的图像，或第一个相似度权重小于10的图像），不过可以自行修改代码实现检索多个结果。

代码：

import os
import shutil
import time
from skimage.metrics import structural_similarity as compare_ssim
from torchvision.transforms import transforms
from PIL import Image
import cv2
import numpy as np

data_transform = transforms.Resize((32, 32))  # 数据库图像和查询图像统一大小，大小为32×32
show_transform = transforms.Resize((224, 224))  # 显示图像大小为224×224


def transformImg(img, transform):
    img = transform(Image.fromarray(img))
    img = np.array(img)
    return img


root_path = "./queryImgs"  # 查询图像所在的文件夹
dataset_path = "./datasetImgs"  # 数据库图像所在的文件夹

for query_img in os.listdir(root_path):
    query_img_path = os.path.join(root_path, query_img)
    query_img_obj = cv2.imread(query_img_path)
    query_img_obj = transformImg(query_img_obj, data_transform)

    best_mse = np.Inf
    best_ssim = np.Inf
    best_weight = np.Inf
    best_img_name = ""
    best_img_path = ""
    best_img_obj = None
    print("Start search Img: ", query_img)
    start_time = time.time()
    for dataset_img in os.listdir(dataset_path):
        dataset_img_path = os.path.join(dataset_path, dataset_img)
        dataset_img_obj = cv2.imread(dataset_img_path)
        # # 统一数据库图像大小, 若数据库图像大小一致则可以只调整查询图像大小。
        # dataset_img_obj = transformImg(dataset_img_obj, data_transform)
        mse = ((query_img_obj - dataset_img_obj) ** 2).mean()
        ssim = compare_ssim(query_img_obj, dataset_img_obj, channel_axis=query_img_obj.shape[2] - 1)
        weight = mse * (1 - ssim)
        if weight < best_weight:
            best_mse = mse
            best_ssim = ssim
            best_weight = weight
            best_img_path = dataset_img_path
            best_img_obj = dataset_img_obj
            best_img_name = dataset_img
            print(query_img, "->", dataset_img, ": ")
            print("\tmse: ", best_mse, " ssim: ", ssim, " weight: ", weight)

        # 权重小于10提前结束检索
        if best_weight < 10:
            break

    elapsed_time = time.time() - start_time

    best_img = np.hstack([transformImg(query_img_obj, show_transform), transformImg(best_img_obj, show_transform)])
    cv2.imshow("left: {}   right: {}".format(query_img, best_img_name), best_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    if not os.path.exists("./dstImgs"): os.mkdir("./dstImgs")
    shutil.copy(best_img_path, './dstImgs/' + best_img_name)
    print("save as: ", './dstImgs/' + best_img_name, " time elapsed: ", elapsed_time, "\n")

结果：