wko_anime-face-similarity/metrics.py

82 lines
3.0 KiB
Python
Raw Normal View History

2023-01-29 22:43:45 +01:00
import cv2
import numpy as np
from skimage.metrics import structural_similarity
2023-01-29 22:43:45 +01:00
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))
2023-01-29 22:57:29 +01:00
def euclidean_distance(data_a: np.ndarray, data_b: np.ndarray) -> float:
gray_a = cv2.cvtColor(data_a, cv2.COLOR_BGR2GRAY)
histogram_a = cv2.calcHist([gray_a], [0], None, [256], [0, 256])
gray_b = cv2.cvtColor(data_b, cv2.COLOR_BGR2GRAY)
histogram_b = cv2.calcHist([gray_b], [0], None, [256], [0, 256])
result, i = [0.], 0
while i < len(histogram_a) and i < len(histogram_b):
result += (histogram_a[i] - histogram_b[i]) ** 2
i += 1
return result[0] ** (1 / 2)
2023-02-01 19:55:12 +01:00
def get_top_results(all_metrics: list[dict], metric='correlation', count=1):
all_metrics.sort(reverse=True, key=lambda item: item['metrics'][metric])
return list(map(lambda item: {'name': item['name'], 'score': item['metrics'][metric]}, all_metrics[:count]))
class AccuracyGatherer:
all_metric_names = [
'structural-similarity',
'euclidean-distance',
'chi-square',
'correlation',
'intersection',
'bhattacharyya-distance'
]
def __init__(self, count, top_ks=(1, 3, 5)):
self.top_ks = top_ks
self.hits = {k: {metric: 0 for metric in AccuracyGatherer.all_metric_names} for k in top_ks}
self.count = count
def print(self):
for k in self.top_ks:
all_metrics = {metric: self.hits[k][metric] / self.count for metric in AccuracyGatherer.all_metric_names}
print(f'Top {k} matches results:')
[print(f'\t{key}: {value * 100}%') for key, value in all_metrics.items()]
def for_results(self, results, test_label):
top_results_all_metrics = {
k: {m: get_top_results(results, m, k) for m in AccuracyGatherer.all_metric_names} for k in self.top_ks
}
for metric_name in AccuracyGatherer.all_metric_names:
self.add_if_hit(top_results_all_metrics, test_label, metric_name)
def add_if_hit(self, results, test_label, metric_name):
for k in self.top_ks:
if any(map(lambda single_result: single_result['name'] == test_label, results[k][metric_name])):
self.hits[k][metric_name] += 1