IQDB-like image similarity calculation in Python

iqdb.py

#!/usr/bin/env python

# Install deps:
#     pip install scikit-image PyWavelets
# Run:
#     python iqdb.py a.jpg b.jpg
# References:
#     https://iqdb.org/code/
#     https://github.com/ricardocabral/iskdaemon
#     https://grail.cs.washington.edu/projects/query/
#     https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/mrquery.pdf

import sys
from heapq import nlargest

from pywt import coeffs_to_array, wavedec2
from skimage.color import rgb2yiq, rgba2rgb
from skimage.io import imread
from skimage.transform import downscale_local_mean, resize

COEFFS = 40
WEIGHTS = {
    "scanned": [
        [5.00, 19.21, 34.37],
        [0.83, 1.26, 0.36],
        [1.01, 0.44, 0.45],
        [0.52, 0.53, 0.14],
        [0.47, 0.28, 0.18],
        [0.30, 0.14, 0.27],
    ],
    "sketch": [
        [4.04, 15.14, 22.62],
        [0.78, 0.92, 0.40],
        [0.46, 0.53, 0.63],
        [0.42, 0.26, 0.25],
        [0.41, 0.14, 0.15],
        [0.32, 0.07, 0.38],
    ],
}
IDX_TO_BIN_MAP = [
    int(min(5, max(i / 16, j / 16)))
    for i in range(128)
    for j in range(128)
]


def get_n_largests(coefficients, n=COEFFS):
    def key(i):
        return abs(coefficients[i])

    return [
        -i if coefficients[i] < 0 else i
        for i in nlargest(n, range(1, len(coefficients)), key=key)
    ]


def calculate_signature_for_channel(channel):
    coeffs, _ = coeffs_to_array(wavedec2(channel, "haar"))

    return coeffs[0][0] / channel.size, get_n_largests(coeffs.flatten())


def get_prescale_factor(w, h):
    if w >= 2048 and h >= 2048:
        return 8
    if w >= 1024 and h >= 1024:
        return 4
    if w >= 512 and h >= 512:
        return 2
    return 1


def quality_resize(image, channels):
    scale = get_prescale_factor(image.shape[0], image.shape[1])

    if scale > 1:
        image = downscale_local_mean(
            image,
            (scale, scale, 1) if channels > 1 else (scale, 1),
        )

    return resize(image, (128, 128, channels), anti_aliasing=True)


def calculate_signature_for_image(image):
    shape = image.shape
    n_dims = len(shape)

    if n_dims == 2:
        return [
            calculate_signature_for_channel(quality_resize(image, 1)[..., 0]),
        ]
    elif n_dims == 3:
        color_channels = shape[-1]

        if color_channels == 4:
            image = rgba2rgb(image, (1, 1, 1))
            color_channels = 3

        image = quality_resize(image, color_channels)
        image = rgb2yiq(image) * 256

        return [
            calculate_signature_for_channel(image[..., i])
            for i in range(color_channels)
        ]

    raise ValueError("Invalid image shape")


def is_grayscale(avg_i, avg_q, threshold=0.006):
    return (abs(avg_i) + abs(avg_q)) < threshold


def reduce_colors(sigs):
    if len(sigs) == 3 and is_grayscale(sigs[1][0], sigs[2][0]):
        return [sigs[0]]
    return sigs


def similarity(a, b, weights=WEIGHTS["scanned"]):
    score, scale = 0, 0

    for channel, ((avg_a, sig_a), (avg_b, sig_b)) in enumerate(zip(a, b)):
        score += weights[0][0] * abs(avg_a - avg_b)
        sig_a = sorted(sig_a, key=abs)
        sig_b = sorted(sig_b, key=abs)
        i, j = 0, 0

        while i < COEFFS or j < COEFFS:
            idx_a = sys.maxsize if i == COEFFS else sig_a[i]
            idx_b = sys.maxsize if j == COEFFS else sig_b[j]
            bin = IDX_TO_BIN_MAP[abs(idx_a if idx_a < idx_b else idx_b)]
            weight = weights[bin][channel]
            scale -= weight

            if idx_a == idx_b:
                score -= weight

            i += idx_a <= idx_b
            j += idx_b <= idx_a

    return score * 100 / scale


def main():
    a = reduce_colors(calculate_signature_for_image(imread(sys.argv[1])))
    b = reduce_colors(calculate_signature_for_image(imread(sys.argv[2])))

    print("similarity %.2f%%" % similarity(a, b))


if __name__ == "__main__":
    main()