diff --git a/comparisons.py b/comparisons.py index 7ff6a72..a4defff 100644 --- a/comparisons.py +++ b/comparisons.py @@ -1,6 +1,6 @@ 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) @@ -23,3 +23,7 @@ def histogram_comparison(data_a: np.ndarray, data_b: np.ndarray) -> dict: '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)) diff --git a/main.py b/main.py index 1fb59c7..8046255 100644 --- a/main.py +++ b/main.py @@ -3,7 +3,7 @@ import cv2 import numpy as np import matplotlib.pyplot as plt -from comparisons import histogram_comparison +from comparisons import histogram_comparison, structural_similarity_index # Allows imports from the style transfer submodule sys.path.append('DCT-Net') @@ -44,6 +44,7 @@ def compare_with_anime_characters(data: np.ndarray) -> int: 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)}') def transfer_to_anime(img: np.ndarray): diff --git a/requirements.txt b/requirements.txt index 45b2c87..33c4fb9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ beautifulsoup4==4.11.1 lxml==4.9.2 opencv-python==4.7.0.68 torch==1.13.1 -matplotlib==3.6.3 \ No newline at end of file +matplotlib==3.6.3 +scikit-image==0.19.3 \ No newline at end of file