diff --git a/color_transfer.py b/color_transfer.py index 377d779..14f4185 100644 --- a/color_transfer.py +++ b/color_transfer.py @@ -3,7 +3,7 @@ from sklearn.cluster import KMeans, MiniBatchKMeans import torch import ast import cv2 -from .utils import EuclideanDistance, ManhattanDistance +from .utils import EuclideanDistance, ManhattanDistance, CosineSimilarity @@ -28,7 +28,8 @@ def SwitchColors(image, detected_colors, target_colors, clustering_model, distan distance_methods = { "Euclidean": EuclideanDistance, - "Manhattan": ManhattanDistance + "Manhattan": ManhattanDistance, + "HSV_Similarity": CosineSimilarity } distance_method = distance_methods.get(distance_method) @@ -54,7 +55,7 @@ class PaletteTransferNode: "target_colors": ("COLOR_LIST",), "color_space": (["RGB", "HSV", "LAB"], {'default': 'RGB'}), "cluster_method": (["Kmeans","Mini batch Kmeans"], {'default': 'Kmeans'}, ), - "distance_method": (["Euclidean", "Manhattan"], {'default': 'Euclidean'}, ) + "distance_method": (["Euclidean", "Manhattan", "CosineSimilarity"], {'default': 'Euclidean'}, ) } } return data_in diff --git a/utils.py b/utils.py index ce099c4..f5a08dd 100644 --- a/utils.py +++ b/utils.py @@ -7,3 +7,19 @@ def EuclideanDistance(detected_colors, target_colors): def ManhattanDistance(detected_colors, target_colors): return np.sum(np.abs(detected_colors - target_colors), axis=1) + + +def CosineSimilarity(detected_colors, target_colors): + return np.dot(detected_colors, target_colors) / (np.linalg.norm(detected_colors) * np.linalg.norm(target_colors)) + + +'''def HSV_Color_Similarity(detected_colors, target_colors): + h1, s1, _ = detected_colors + h2, s2, _ = target_colors + + h1_rad, h2_rad = np.radians(h1), np.radians(h2) + + v1 = np.array([s1 * np.cos(h1_rad), s1 * np.sin(h1_rad)]) + v2 = np.array([s2 * np.cos(h2_rad), s2 * np.sin(h2_rad)]) + + return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))''' \ No newline at end of file