当前位置：首页 > news >正文

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

news 2025/7/21 14:58:24

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

背景：由于需要查找别人代码保存的图像的命名，但由于数据集是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 npdata_transform = transforms.Resize((32, 32))  # 数据库图像和查询图像统一大小，大小为32×32
show_transform = transforms.Resize((224, 224))  # 显示图像大小为224×224def transformImg(img, transform):img = transform(Image.fromarray(img))img = np.array(img)return imgroot_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.Infbest_ssim = np.Infbest_weight = np.Infbest_img_name = ""best_img_path = ""best_img_obj = Noneprint("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 = msebest_ssim = ssimbest_weight = weightbest_img_path = dataset_img_pathbest_img_obj = dataset_img_objbest_img_name = dataset_imgprint(query_img, "->", dataset_img, ": ")print("\tmse: ", best_mse, " ssim: ", ssim, " weight: ", weight)# 权重小于10提前结束检索if best_weight < 10:breakelapsed_time = time.time() - start_timebest_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")

结果：