From 752ad34a4e624841ac6a8d0ddbd144f771d03549 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 11 Feb 2025 20:36:50 +0200 Subject: [PATCH] hsv color space --- color_transfer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/color_transfer.py b/color_transfer.py index dcac9ca..f304573 100644 --- a/color_transfer.py +++ b/color_transfer.py @@ -2,6 +2,7 @@ import numpy as np from sklearn.cluster import KMeans, MiniBatchKMeans import torch import ast +import cv2 from .utils import EuclideanDistance, ManhattanDistance @@ -73,6 +74,9 @@ class PaletteTransferNode: for image in image: img = 255. * image.cpu().numpy() + if color_space == "HSV": + img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV) + clustered_img, detected_colors, clustering_model = ColorClustering(img, len(target_colors), cluster_method) processed = SwitchColors(clustered_img, detected_colors, target_colors, clustering_model, distance_method) processedImages.append(processed)