Added color similarity for HSV

color_transfer.py

@@ -3,7 +3,7 @@ from sklearn.cluster import KMeans, MiniBatchKMeans
 import torch
 import ast
 import cv2
-from .utils import EuclideanDistance, ManhattanDistance, CosineSimilarity
+from .utils import EuclideanDistance, ManhattanDistance, CosineSimilarity, HSV_Color_Similarity, Blur
 
 
 
@@ -29,7 +29,8 @@ def SwitchColors(image, detected_colors, target_colors, clustering_model, distan
     distance_methods = {
     "Euclidean": EuclideanDistance,
     "Manhattan": ManhattanDistance,
-    "Cosine Similarity": CosineSimilarity
+    "Cosine Similarity": CosineSimilarity,
+    "HSV Distance": HSV_Color_Similarity
     }
 
     distance_method = distance_methods.get(distance_method)
@@ -55,11 +56,12 @@ 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", "Cosine Similarity"], {'default': 'Euclidean'}, )
+                "distance_method": (["Euclidean", "Manhattan", "Cosine Similarity", "HSV Distance"], {'default': 'Euclidean'}, ),
+                "gaussian_blur": ("BOOLEAN", {'default': False}),
                 }
             }
         return data_in
-
+    """("INT", {"default": 2100})"""
 
     CATEGORY = "Color Transfer"
     RETURN_TYPES = ("IMAGE",)
@@ -67,7 +69,7 @@ class PaletteTransferNode:
     CATEGORY = "Palette Transfer"
 
 
-    def color_transfer(self, image, target_colors, color_space, cluster_method, distance_method):
+    def color_transfer(self, image, target_colors, color_space, cluster_method, distance_method, gaussian_blur):
 
         if len(target_colors) == 0:
             return (image,)
@@ -91,6 +93,9 @@ class PaletteTransferNode:
             if color_space == "HSV":
                 processed = cv2.cvtColor(processed, cv2.COLOR_HSV2RGB)
             
+            if gaussian_blur:
+                processed = Blur(processed, 3)
+
             processed = np.array(processed).astype(np.float32) / 255.0
             processedImage = torch.from_numpy(processed)[None,]
 

utils.py

@@ -1,4 +1,5 @@
 import numpy as np
+import cv2
 
 
 def EuclideanDistance(detected_color, target_colors):
@@ -13,13 +14,35 @@ def CosineSimilarity(detected_color, target_colors):
     return -np.dot(target_colors, detected_color) / (np.linalg.norm(detected_color) * np.linalg.norm(target_colors, axis=1))
 
 
-'''def HSV_Color_Similarity(detected_colors, target_colors):
-    h1, s1, _ = detected_colors
-    h2, s2, _ = target_colors
+def HSV_Color_Similarity(detected_color, target_colors):
+    detected_color = np.array(detected_color)
+    target_colors = np.array(target_colors)
     
-    h1_rad, h2_rad = np.radians(h1), np.radians(h2)
+    h1, s1, _ = detected_color
+    h2 = target_colors[:, 0]
+    s2 = target_colors[:, 1]
+
+    h1_rad = np.radians(h1)
+    h2_rad = 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)])
+    v1_x = s1 * np.cos(h1_rad)
+    v1_y = s1 * np.sin(h1_rad)
+    v1 = np.array([v1_x, v1_y])
     
-    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))'''
+
+    v2_x = s2 * np.cos(h2_rad)
+    v2_y = s2 * np.sin(h2_rad)
+    v2 = np.vstack([v2_x, v2_y])
+    
+    dot_products = np.dot(v1, v2)
+    
+    v1_norm = np.linalg.norm(v1)
+    v2_norms = np.linalg.norm(v2, axis=0)
+    
+    similarities = dot_products / (v1_norm * v2_norms)
+    
+    return -similarities
+
+
+def Blur(image, kernel_size):
+    return cv2.medianBlur(image.astype(np.uint8), kernel_size)