Простой поиск дубликатов изображения

Поиск копий изображения, разбросанных по разным директориям хранилища, одна из частых проблем пользователя компьютера. Существует множество проверенных решений, основанных на разных алгоритмах для решения проблемы. Этот пример использует элементы машинного обучения, текущий уровень развития инструментов, позволяет с минимальными усилиями решать "бытовые задачи", к которым относиться поиск и удаление копий. В качестве меры сходства - косинусное сходство. Сравнение многомерных массивов (изображение в цифровом пространстве), ресурсоемкий процесс, поэтому, применяем обученную свёрточную нейронную сеть для уменьшения размерности с учетом важных пространственных признаков. Библиотека keras содержит готовые модели под разные задачи, этот пример задействует архитектуру VGG16 обученную на данных imagenet. Вход в сеть (N, 224, 224, 3), выход (1, 512).

Примерный алгоритм

1. Получить изображения из нужных папок

2. Преобразовать изображение в уменьшенный вектор с помощью CNN

3. Сравнить векторы между собой

4. Сохранить результаты

Полезные инструменты в Python

• numpy 1.19.5

• opencv 4.5.1

• keras 2.1.5

# -*- coding: utf-8 -*-
import sys, os, time
import numpy as np
import keras
from keras.preprocessing import image as image_utils
import json

# самый простой парсинг файлов
class DATA():
    def __init__(self):
        self.file = []
    def parseIMG(self, dir_name):
        path = f"{dir_name}/"
        for r, d, f in os.walk(path):
            for ix, file in enumerate(f): 
                       if ".png" in file.lower(): 
                           self.file.append(os.path.join(r, file))
                       elif ".jpg" in file.lower(): 
                           self.file.append(os.path.join(r, file))
                       elif ".jpeg" in file.lower(): 
                           self.file.append(os.path.join(r, file))

# преобразуем многомерную матрицу изображения в вектор 1x512 
def deep_vector(x):
       t_arr = image_utils.load_img(x, target_size=(224, 224))
       t_arr = image_utils.img_to_array(t_arr)
       t_arr = np.expand_dims(t_arr, axis=0)
       processed_img = preprocess(t_arr)
       preds = model.predict(processed_img)
       return preds                          
                          
                          
# косинусное сходство
def similarity(vector1, vector2):
        return np.dot(vector1, vector2.T) / np.dot(np.linalg.norm(vector1, axis=1, keepdims=True), 
                                                   np.linalg.norm(vector2.T, axis=0, keepdims=True))

# очень простой алгоритм сортировки
def func_sort(ID):
    while len(arr_list) != 0:
        for ix, i in enumerate(arr_list):
            G[ID] = [arr.file[ix]]
            del arr.file[ix]
            del arr_list[ix]
            for iix, ii in enumerate(arr_list):
                    KEF = similarity(i, ii) 
                    KEF = float(KEF[0][0])
                    if KEF > tresh:
                        G[ID].append(arr.file[iix])
                        del arr.file[iix]
                        del arr_list[iix] 
            ID += 1

if __name__ == '__main__':
    arr = DATA() 
  	arr.parseIMG(sys.argv[1]) # путь к директории

    # обученная сеть VGG16 на данных imagenet
    model = keras.applications.vgg16.VGG16(include_top=False, 
                                           weights='imagenet', 
                                           input_tensor=None, 
                                           input_shape=None, 
                                           pooling='max')
    # функция создания пакета изображений 
    # в этом примере обрабатываем по одному изображению
    preprocess = keras.applications.vgg16.preprocess_input

    tresh = .9998 # пороговое значение 
    G = {} # хранение полученных результатов
    arr_list = [] # временное хранилище векторов изображения
    error_list = [] # ошибки
		
    # создаем вектор изображения
    # помещаем в списко arr_list
    for i in arr.file:
      	try:
            _vector = deep_vector(i)
            arr_list.append(_vector)
        except Exception as e:
          	error_list.append(i) 
    
    # функция сортировки
    func_sort(0)
    
    # сохраняем полученное 
    with open('data.json', 'w') as _file:
        json.dump(G, fp)

    with open('data_error.json', 'w') as _file:
        json.dump(error_list, fp)   
  

ImageNet - содержит большое количество классов, сети обученные на этих данных, хорошо справляются с поиском дубликатов. Можно поэкспериментировать с пороговым значением и получить простой инструмент для поиска похожих изображений, но для таких систем рекомендуется обучать сеть на своих данных с использованием потерь triplet или arcface.