CV系列之用图像相似度去重

当我们处理一批图片时，这批图片可能会有重复的图片，下面分享一个具备图片相似度去重功能的脚本，欢迎大家对该脚本进行改进。

1. import os
   
2. import cv2
   
3. from PIL import Image
   
4. import imagehash
   
5. import shutil
   
6. import numpy as np
   
7. from tqdm import tqdm
   
8. 
   
9. def calculate_phash(image):
   
10.     # 使用感知哈希算法计算图像的哈希值
    
11.     pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    
12.     phash = imagehash.phash(pil_image)
    
13.     return str(phash)
    
14. 
    
15. def calculate_mse(image1, image2):
    
16.     # 计算均方误差
    
17.     mse = ((image1 - image2) ** 2).mean()
    
18.     return mse
    
19. 
    
20. def move_similar_images(source_folder, destination_folder, near_num=10, mse_thred = 100):
    
21.     # 获取源文件夹中所有图片的路径
    
22.     image_paths = [os.path.join(source_folder, image_name) for image_name in os.listdir(source_folder)]
    
23. 
    
24.     # 按照图片名称排序
    
25.     image_paths.sort()
    
26. 
    
27.     # 遍历图片列表，判断相邻的图片是否相似
    
28.     img_num = len(image_paths)-near_num
    
29.     progress_bar = tqdm(total=img_num, unit='files', desc='processing files')
    
30.     for i in range(len(image_paths)-near_num):
    
31.         progress_bar.set_postfix(file=i)
    
32.         progress_bar.update(1)
    
33. 
    
34.         # current_image = cv2.imread(image_paths)
    
35.         current_image = cv2.imdecode(np.fromfile(image_paths[i], dtype=np.uint8), -1)
    
36. 
    
37.         # 比较当前图像和后续 near_num 张图像
    
38.         for j in range(i+1, i+near_num+1):
    
39.             # next_image = cv2.imread(image_paths[j])
    
40.             next_image = cv2.imdecode(np.fromfile(image_paths[j], dtype=np.uint8), -1)
    
41. 
    
42.             #
    
43.             if current_image.shape != next_image.shape:
    
44.                 continue      
    
45. 
    
46.             # # 计算当前图像和下一张图像的哈希值，如果相同，移动
    
47.             # current_hash = calculate_phash(current_image)
    
48.             # next_hash = calculate_phash(next_image)
    
49.             # if current_hash == next_hash:
    
50.             #     destination_path = os.path.join(destination_folder, os.path.basename(image_paths[i]))
    
51.             #     shutil.move(image_paths[i], destination_path)
    
52.             #     break
    
53.         
    
54.             # 计算mse相似度，如果低于阈值，移动
    
55.             try:
    
56.                 mse = calculate_mse(current_image, next_image)
    
57.             except IOError as e:
    
58.                 print('erro:', e)            
    
59.             if mse < mse_thred:
    
60.                 destination_path = os.path.join(destination_folder, os.path.basename(image_paths[i]))
    
61.                 shutil.move(image_paths[i], destination_path)
    
62.                 break
    
63. 
    
64.     progress_bar.close()   
    
65. 
    
66. 
    
67. if __name__ == '__main__':
    
68.     # 源文件夹和目标文件夹的路径
    
69.     src_dir = r'F:\data\xxx'
    
70.     des_dir = src_dir + '_chongfu'
    
71.     if not os.path.exists(des_dir):
    
72.         os.makedirs(des_dir)
    
73. 
    
74.     # 移动相似图片到目标文件夹
    
75.     move_similar_images(src_dir, des_dir, near_num=20, mse_thred = 70)

复制代码

除了计算哈希值，也可以通过一些网络模型提取特征值来进行比对，相似度还可以用余弦相似度，计算向量内积就可以