similar_image_finder.py

from datetime import timedelta
from glob import glob
from os import mkdir
from os.path import exists
from shutil import move
from time import time
import cv2
from colorama import init, Style, Fore

bright = Style.BRIGHT
green, blue, red, cyan, reset = Fore.GREEN + bright, Fore.BLUE + bright, Fore.RED + bright, Fore.CYAN, Fore.RESET
init(convert=True, autoreset=True)


def img_similarity(file):
    img_compare = cv2.imread(file)
    sift = cv2.xfeatures2d.SIFT_create()
    _, desc_1 = sift.detectAndCompute(ref_img, None)
    _, desc_2 = sift.detectAndCompute(img_compare, None)
    index_params = dict(algorithm=0, trees=5)
    search_params = dict(checks=search_times)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    matches = flann.knnMatch(desc_1, desc_2, k=2)
    good_points = []
    [good_points.append(m) for m, n in matches if m.distance < 0.7 * n.distance]
    [move(file, "matched") if len(good_points) >= min_match_points else None]


def folder_list():
    folder = input('[+] Drag and drop the folder :- ').replace('"', '')
    img = glob(f'{folder}/**/*.jpg', recursive=True) + (glob(f'{folder}/**/*.png', recursive=True)) + (glob(f'{folder}/**/*.jpeg', recursive=True))
    print(f"{green}[+] Total files in folder are {len(img)}")
    for img_file in range(len(img)):
        print(f'{blue}[{img_file + 1}] {green}processing the image {cyan}{img[img_file]}')
        try:
            img_similarity(img[img_file])
        except Exception as e:
            print(f"{red}[-] Error occurred while processing image {img[img_file]} :- {e}")


def main():
    global ref_img, min_match_points, search_times
    [mkdir(folder) for folder in ["matched"] if not exists(folder)]
    ref_img = cv2.imread(input("[+] Enter main image for reference :- ").replace('"', ""))
    min_match_points = int(input("[+] Enter minimum match points(Default 100) :-") or 100)
    search_times = int(input("[+] Enter the number os times to search for images higher the value higher accuracy or press enter to use default(Default 0) :-") or 0)
    folder_list()
    time_output = time() - start_time
    sec = str(timedelta(seconds=(int(time_output)))).split(":")
    print(f"[*]{green} Time taken to process images is -{reset} {sec[0]} H : {sec[1]} M : {sec[2]} S")
    print(f"{green}[~] Successfully completed [~]")


def credit():
    credit_text = f"""
               {red}SIMILAR IMAGE FINDER USING BASE IMAGE
{green}               
     ██▀███   ▄▄▄       ██░ ██  █    ██  ██▓    
    ▓██ ▒ ██▒▒████▄    ▓██░ ██▒ ██  ▓██▒▓██▒    
    ▓██ ░▄█ ▒▒██  ▀█▄  ▒██▀▀██░▓██  ▒██░▒██░    
    ▒██▀▀█▄  ░██▄▄▄▄██ ░▓█ ░██ ▓▓█  ░██░▒██░    
    ░██▓ ▒██▒ ▓█   ▓██▒░▓█▒░██▓▒▒█████▓ ░██████▒
    ░ ▒▓ ░▒▓░ ▒▒   ▓▒█░ ▒ ░░▒░▒░▒▓▒ ▒ ▒ ░ ▒░▓  ░
      ░▒ ░ ▒░  ▒   ▒▒ ░ ▒ ░▒░ ░░░▒░ ░ ░ ░ ░ ▒  ░
      ░░   ░   ░   ▒    ░  ░░ ░ ░░░ ░ ░   ░ ░   
       ░           ░  ░ ░  ░  ░   ░         ░  ░ {blue}code generated by Rahul.p\n
    """
    print(credit_text)


try:
    start_time = time()
    credit()
    main()
except KeyboardInterrupt:
    print(f'{red}\n[~] Exiting ....')
except Exception as e:
    print(f"{red}\n[-] error message is {e}")