SquirrelDetector/scripts/dedup.py

import argparse
import glob
import json
import numpy as np
import os
import sys
import torch

from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity
from torchvision import transforms
from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
from torchvision.models.feature_extraction import create_feature_extractor


class FeatureExtractor(torch.nn.Module):

    def __init__(self):
        super(FeatureExtractor, self).__init__()

        model = efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1)
        return_nodes = {'flatten': 'feature_vector'}
        self._feature_extractor = create_feature_extractor(model, return_nodes)
        self._feature_extractor.eval()

        self._preprocess = transforms.Compose([
            transforms.Resize(224), 
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def forward(self, img):
        data = Image.open(img)
        data = data.convert('RGB')
        x = self._preprocess(data).unsqueeze(0)
        with torch.no_grad():
            y = self._feature_extractor(x)['feature_vector']
            return y.squeeze().numpy()
         

def is_image(filename):
    ext = os.path.splitext(filename)[-1]
    return ext.lower() in ('.jpg', '.jpeg', '.png')


def extract_features(images):
    extractor = FeatureExtractor()
    features = []
    for image in images:
        features.append(extractor(image))
    return np.array(features)        


def main():
    parser = argparse.ArgumentParser('Creates a list of duplicate images found within a directory and saves it to a file (similarity.json)')
    parser.add_argument('input', help='Input directory with images to compare')
    parser.add_argument('--threshold', type=float, default=0.95, help='Threshold to determine similarity')

    args = parser.parse_args()
    threshold = args.threshold

    images = [f for f in glob.glob(os.path.join(args.input, '*')) if is_image(f)]
    img_names = np.array([os.path.split(i)[-1] for i in images])

    print(f'Anylizing {len(images)} images...')
    
    print(f'    + Extracting features...')
    features = extract_features(images)

    print(f'    + Calculating cosine simularity...')
    scores = cosine_similarity(features)
    np.fill_diagonal(scores, 0.0)

    print(f'    + Computing output')
    similarity_set = set()
    for i, j in enumerate(scores):
        pred = (j >= threshold) 
        same = img_names[pred].tolist()
        if len(same):
            same.append(img_names[i])
            same.sort()
            similarity_set.add(';'.join(same))

    similarity_list = [s.split(';') for s in sorted(list(similarity_set))]

    with open('similarity.json', mode='w') as f:
        json.dump(similarity_list, f) 
    
if __name__ == '__main__':
    main()