From 49e337e5e97625fc91c944ed7ace675ee894f800 Mon Sep 17 00:00:00 2001
From: Marcin Kostrzewski <komar86311@gmail.com>
Date: Tue, 31 Jan 2023 20:48:24 +0100
Subject: [PATCH] Compare against the entire anime faces dataset

---
 load_test_data.py |  1 -
 main.py           | 39 ++++++++++++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/load_test_data.py b/load_test_data.py
index b5a635b..300cd25 100644
--- a/load_test_data.py
+++ b/load_test_data.py
@@ -29,7 +29,6 @@ def load_data(input_dir, newSize=(64,64)):
         p = image_path / n
         img = imread(p) # zwraca ndarry postaci xSize x ySize x colorDepth
         img = cv.resize(img, newSize, interpolation=cv.INTER_AREA) # zwraca ndarray
-        img = img / 255 # type: ignore #normalizacja
         test_img.append(img)
         labels.append(n)
 
diff --git a/main.py b/main.py
index af5ad1b..e53c97b 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 from comparisons import histogram_comparison, structural_similarity_index, euclidean_distance
+from load_test_data import load_data
 
 # Allows imports from the style transfer submodule
 sys.path.append('DCT-Net')
@@ -36,16 +37,31 @@ def plot_two_images(a: np.ndarray, b: np.ndarray):
   plt.show()
 
 
-def compare_with_anime_characters(data: np.ndarray) -> int:
-  # Example will be one face from anime dataset
-  example = load_source('data/images/Aisaka, Taiga.jpg')
-  # TODO: Use a different face detection method for anime images
-  example_face = find_and_crop_face(example, 'haarcascades/lbpcascade_animeface.xml')
-  data_rescaled = cv2.resize(data, example_face.shape[:2])
-  plot_two_images(example_face, data_rescaled)
-  print(histogram_comparison(data_rescaled, example_face))
-  print(f'structural-similarity: {structural_similarity_index(data_rescaled, example_face)}')
-  print(f'euclidean-distance: {euclidean_distance(data_rescaled, example_face)}')
+def compare_with_anime_characters(source: np.ndarray, verbose=False) -> list[dict]:
+  dataset = load_data('data/images')
+  all_metrics = []
+  for anime_image, label in zip(dataset['values'], dataset['labels']):
+    current_result = {
+      'name': label,
+      'metrics': {}
+    }
+    # TODO: Use a different face detection method for anime images
+    # anime_face = find_and_crop_face(anime_image, 'haarcascades/lbpcascade_animeface.xml')
+    anime_face = anime_image
+    source_rescaled = cv2.resize(source, anime_face.shape[:2])
+    if verbose:
+      plot_two_images(anime_face, source_rescaled)
+    current_result['metrics'] = histogram_comparison(source_rescaled, anime_face)
+    current_result['metrics']['structural-similarity'] = structural_similarity_index(source_rescaled, anime_face)
+    current_result['metrics']['euclidean-distance'] = euclidean_distance(source_rescaled, anime_face)
+    all_metrics.append(current_result)
+
+  return all_metrics
+
+
+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]))
 
 
 def transfer_to_anime(img: np.ndarray):
@@ -57,4 +73,5 @@ if __name__ == '__main__':
   source = load_source('UAM-Andre.jpg')
   source_anime = transfer_to_anime(source)
   source_face_anime = find_and_crop_face(source_anime)
-  print(compare_with_anime_characters(source_face_anime))
+  results = compare_with_anime_characters(source_face_anime)
+  print(get_top_results(results, count=5))