wko_anime-face-similarity/comparisons.py

import cv2
import numpy as np
from skimage.metrics import structural_similarity

def histogram_comparison(data_a: np.ndarray, data_b: np.ndarray) -> dict:
  hsv_a = cv2.cvtColor(data_a, cv2.COLOR_BGR2HSV)
  hsv_b = cv2.cvtColor(data_b, cv2.COLOR_BGR2HSV)

  histSize = [50, 60]
  hue_ranges = [0, 180]
  sat_ranges = [0, 256]
  channels = [0, 1]
  ranges = hue_ranges + sat_ranges

  hist_a = cv2.calcHist([hsv_a], channels, None, histSize, ranges, accumulate=False)
  cv2.normalize(hist_a, hist_a, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)
  hist_b = cv2.calcHist([hsv_b], channels, None, histSize, ranges, accumulate=False)
  cv2.normalize(hist_b, hist_b, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)

  return {
    'correlation': cv2.compareHist(hist_a, hist_b, 0),
    'chi-square': cv2.compareHist(hist_a, hist_b, 1),
    'intersection': cv2.compareHist(hist_a, hist_b, 2),
    'bhattacharyya-distance': cv2.compareHist(hist_a, hist_b, 3),
  }


def structural_similarity_index(data_a: np.ndarray, data_b: np.ndarray) -> float:
  return structural_similarity(cv2.cvtColor(data_a, cv2.COLOR_BGR2GRAY), cv2.cvtColor(data_b, cv2.COLOR_BGR2GRAY))