find_best_img_match.py

from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from tqdm import tqdm
import cv2
import matplotlib.pyplot as plt
import numpy as np


def find_best_img_match(target_img: np.array, imgpaths: list[Path]):
    """
    Find the image in imagepaths with the best correlation with target_img
    """
    best_corr = 0.0
    best_path = None
    target_ = target_img.astype(np.float32)
    assert len(target_.shape) == 2, "Must use grayscale"
    h, w = target_.shape

    def task(imgpath):
        img = cv2.imread(str(imgpath), cv2.IMREAD_GRAYSCALE).astype(np.float32)
        _, corr = cv2.phaseCorrelate(target_, cv2.resize(img, (w, h)))
        return corr, imgpath

    with ThreadPoolExecutor(max_workers=8) as executor:
        for corr, imgpath in tqdm(executor.map(task, imgpaths), total=len(imgpaths)):
            if corr > best_corr:
                best_corr, best_path = corr, imgpath
                print(f"New best match {corr=} {imgpath=}")

    print(f"Final best match {best_corr=} {best_path=}")
    return best_corr, best_path