筛查重复或类似图像

筛查重复或类似图像

import subprocess
import sys
import os
import cv2
import numpy as np
from keras.applications.resnet50 import ResNet50, preprocess_input
from sklearn.metrics.pairwise import cosine_similarity  # 用于余弦相似度

# pip 安装 scikit-learn
def install_package(package):
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
        print(f"{package} 安装成功！")
    except subprocess.CalledProcessError as e:
        print(f"安装失败: {e}")

# 确保安装 scikit-learn
install_package("scikit-learn")

# 提取图片特征
def extract_image_features(image_path):
    if not os.path.exists(image_path):  # 检查文件是否存在
        print(f"文件不存在: {image_path}")
        return None
    
    image = cv2.imread(image_path)  # 读取图片
    if image is None:  # 检查图片是否加载成功
        print(f"无法读取图片: {image_path}")
        return None  # 如果读取失败，返回None
    image = cv2.resize(image, (256, 256))  # 缩放图片到统一尺寸
    image = image[16:240, 16:240]  # 裁剪中间区域(224x224)
    
    image = np.expand_dims(image, axis=0)  # 扩展维度以匹配模型输入要求
    image = preprocess_input(image)  # 预处理图片
    
    features = model.predict(image)  # 提取特征向量
    features /= np.linalg.norm(features)  # 归一化特征向量
    
    print(f"特征向量: {features.flatten()}")  # 打印特征向量，看看是否有差异
    return features.flatten()  # 平铺特征向量

# 移动重复图片到指定文件夹
def move_duplicate_images(directory, output_directory, use_cosine_similarity=True):
    current_dir = directory  # 使用传入的目录路径
    files = [f for f in os.listdir(current_dir) if os.path.isfile(os.path.join(current_dir, f))]  # 获取目录下的所有文件
    
    image_features = {}
    moved_count = 0  # 记录移动的图片数量
    duplicate_pairs = []  # 用于保存重复图片的文件名对

    # 确保输出目录存在
    if not os.path.exists(output_directory):  # 如果目录不存在，创建它
        os.makedirs(output_directory)

    for file_name in files:
        if file_name.endswith(".jpg") or file_name.endswith(".png"):  # 筛选出图片文件
            file_path = os.path.join(current_dir, file_name)
            image_feature = extract_image_features(file_path)

            if image_feature is None:  # 如果读取失败，跳过此文件
                continue

            is_duplicate = False
            for existing_path, existing_feature in image_features.items():
                if use_cosine_similarity:
                    # 使用余弦相似度计算相似度
                    similarity = cosine_similarity([existing_feature], [image_feature])[0][0]
                    print(f"计算相似度: {similarity}")  # 打印计算出的相似度

                    if similarity > 0.8:  # 设置相似度阈值，值越接近1说明越相似
                        is_duplicate = True
                        print(f"移动重复图片: {file_path}")
                        # 移动文件到指定文件夹
                        new_path = os.path.join(output_directory, file_name)
                        os.rename(file_path, new_path)  # 移动文件
                        moved_count += 1
                        # 记录重复的文件名对 (当前文件名和已有文件名)
                        duplicate_pairs.append((file_name, os.path.basename(existing_path)))
                        break
                else:
                    # 使用欧氏距离计算相似度
                    distance = np.linalg.norm(existing_feature - image_feature)  # 计算欧氏距离
                    print(f"计算距离: {distance}")  # 打印计算出的欧氏距离
                    if distance < 0.6:  # 删除类似图像 0.6  删除重复图像 0.1
                        is_duplicate = True
                        print(f"移动重复图片: {file_path}")
                        # 移动文件到指定文件夹
                        new_path = os.path.join(output_directory, file_name)
                        os.rename(file_path, new_path)  # 移动文件
                        moved_count += 1
                        # 记录重复的文件名对 (当前文件名和已有文件名)
                        duplicate_pairs.append((file_name, os.path.basename(existing_path)))
                        break

            if not is_duplicate:
                image_features[file_path] = image_feature

    # 将重复图片文件名对保存到txt文件
    output_file = os.path.join(output_directory, "duplicate_images.txt")  # 设置输出文件路径
    if duplicate_pairs:
        with open(output_file, "w") as f:
            for file1, file2 in duplicate_pairs:
                f.write(f"{file1} 与 {file2} 重复或类似\n")

    print("已移动 {} 张重复图片".format(moved_count))

# 加载预训练的ResNet50模型
model = ResNet50(weights='imagenet', include_top=False, pooling='avg')

# 指定目录路径，设置是否使用余弦相似度
input_directory = r"D:\test\img"  # 图片所在文件夹
output_directory = r"D:\test\img2"  # 自定义保存重复图片的文件夹

move_duplicate_images(input_directory, output_directory, use_cosine_similarity=False)  # 设置为True使用余弦相似度判断两个图片的特征方向是否相似，设置为False移除重复图像