From 5fa20b44e9a889c8e0f6614ad73c63b9ff8d37ef Mon Sep 17 00:00:00 2001
From: unknown <alexyalynski@gmail.com>
Date: Tue, 11 Feb 2025 20:21:16 +0200
Subject: [PATCH] new params

---
 color_transfer.py | 25 ++++++++++---------------
 utils.py          |  9 +++++++++
 2 files changed, 19 insertions(+), 15 deletions(-)
 create mode 100644 utils.py

diff --git a/color_transfer.py b/color_transfer.py
index 48a6aaa..dcac9ca 100644
--- a/color_transfer.py
+++ b/color_transfer.py
@@ -2,14 +2,7 @@ import numpy as np
 from sklearn.cluster import KMeans, MiniBatchKMeans
 import torch
 import ast
-
-
-def EuclideanDistance(detected_colors, target_colors):
-    return np.linalg.norm(detected_colors - target_colors, axis=1)
-
-
-def ManhattanDistance(detected_colors, target_colors):
-    return np.sum(np.abs(detected_colors - target_colors), axis=1)
+from .utils import EuclideanDistance, ManhattanDistance
 
 
 def ColorClustering(image, k, cluster_method):
@@ -57,7 +50,8 @@ class PaletteTransferNode:
         data_in = {
             "required": {
                 "image": ("IMAGE",),
-                "target_colors": ("COLORS",),
+                "target_colors": ("COLOR_LIST",),
+                "color_space": ("COLOR_SPACE",),
                 "cluster_method": (["Kmeans","Mini batch Kmeans"], {'default': 'Kmeans'}, ),
                 "distance_method": (["Euclidean", "Manhattan"], {'default': 'Euclidean'}, )
                 }
@@ -69,7 +63,7 @@ class PaletteTransferNode:
     CATEGORY = "Palette Transfer"
 
 
-    def color_transfer(self, image, target_colors, cluster_method, distance_method):
+    def color_transfer(self, image, target_colors, color_space, cluster_method, distance_method):
 
         if len(target_colors) == 0:
             return (image,)
@@ -93,13 +87,14 @@ class ColorPaletteNode:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "color_palette": ("STRING", {'default': '', 'multiline': True})
+                "color_palette": ("STRING", {'default': '[(30, 32, 30), (60, 61, 55), (105, 117, 101), (236, 223, 204)]', 'multiline': True}),
+                "color_space": (["RGB", "HSV", "LAB"], {'default': 'RGB'}),
             },
         }
 
-    RETURN_TYPES = ("COLORS", )
-    RETURN_NAMES = ("Color palette", )
+    RETURN_TYPES = ("COLOR_LIST", "COLOR_SPACE")
+    RETURN_NAMES = ("Color palette", "Color space")
     FUNCTION = "color_list"
 
-    def color_list(self, color_palette):
-        return (ast.literal_eval(color_palette), )
+    def color_list(self, color_palette, color_space):
+        return (ast.literal_eval(color_palette), color_space, )
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..ce099c4
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,9 @@
+import numpy as np
+
+
+def EuclideanDistance(detected_colors, target_colors):
+    return np.linalg.norm(detected_colors - target_colors, axis=1)
+
+
+def ManhattanDistance(detected_colors, target_colors):
+    return np.sum(np.abs(detected_colors - target_colors), axis=1)