From 2f147854a655bd4ec54018e60b0ab9798fc4a5e6 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Wed, 7 Aug 2024 00:33:54 +0200
Subject: [PATCH] Add colorspace conversion nodes

---
 __init__.py         |  4 ++++
 nodes/colorspace.py | 48 +++++++++++++++++++++++++++++++++++++++++++++
 nodes/tonemap.py    |  2 +-
 pyproject.toml      | 14 +++++++++++++
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 nodes/colorspace.py
 create mode 100644 pyproject.toml

diff --git a/__init__.py b/__init__.py
index a29a57b..93f13ee 100644
--- a/__init__.py
+++ b/__init__.py
@@ -39,6 +39,10 @@ else:
 		from .nodes.tonemap import NODE_CLASS_MAPPINGS as tonemap_nodes
 		NODE_CLASS_MAPPINGS.update(tonemap_nodes)
 
+		# color space nodes
+		from .nodes.colorspace import NODE_CLASS_MAPPINGS as colorspace_nodes
+		NODE_CLASS_MAPPINGS.update(colorspace_nodes)
+
 	# export
 	NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}
 	__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
diff --git a/nodes/colorspace.py b/nodes/colorspace.py
new file mode 100644
index 0000000..15ae880
--- /dev/null
+++ b/nodes/colorspace.py
@@ -0,0 +1,48 @@
+import cv2
+import torch
+import numpy as np
+
+# You can add any color spaces supported by cv2
+# [x[6:] for x in dir(cv2) if x.startswith("COLOR_")]
+common = [
+	"RGB", "BGR", "HLS", "HSV", "YCrCb", "YUV"
+]
+
+class ColorspaceConvert:
+	def __init__(self):
+		pass
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"image": ("IMAGE",),
+				"src": (common, {"default": "RGB"}),
+				"dst": (common, {"default": "RGB"}),
+			}
+		}
+
+	RETURN_TYPES = ("IMAGE",)
+	FUNCTION = "convert"
+	CATEGORY = "ColorMod"
+	TITLE = "Convert color space"
+
+	def convert(self, image, src, dst):
+		if src == dst:
+			return (image, )
+		
+		atr = getattr(cv2, f"COLOR_{src}2{dst}", None)
+		assert atr, f"Color conversion failed! Missing cv2 op 'COLOR_{src}2{dst}'"
+		
+		out = []
+		for batch in image:
+			img = (batch.cpu().numpy() * 255.0).astype(np.uint8)
+			mod = cv2.cvtColor(img, atr)
+			out.append(torch.from_numpy(mod.copy()) / 255.0)
+		out = torch.stack(out, dim=0)
+		print(torch.min(out), torch.mean(out), torch.max(out))
+		return (out,)
+
+NODE_CLASS_MAPPINGS = {
+	"ColorspaceConvert": ColorspaceConvert,
+}
diff --git a/nodes/tonemap.py b/nodes/tonemap.py
index 5aaf376..77edaf8 100644
--- a/nodes/tonemap.py
+++ b/nodes/tonemap.py
@@ -126,5 +126,5 @@ NODE_CLASS_MAPPINGS = {
 # not all cv2 versions support all tonemap nodes (e.g Durand require nonfree)
 for name in list(NODE_CLASS_MAPPINGS.keys()):
 	if NODE_CLASS_MAPPINGS[name].tonemap_op is None:
-		print(f"Ignoring node '{name}' due to cv2 edition/version")
+		print(f"ColorMod: Ignoring node '{name}' due to cv2 edition/version")
 		del NODE_CLASS_MAPPINGS[name]
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..bdfd3d1
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,14 @@
+[project]
+name = "comfyui_colormod"
+description = "Nodes for contrast/color editing as well as 16bit/HDR/tonemapping support."
+version = "1.2.0"
+license = { file = "LICENSE" }
+dependencies = ["pypng", "opencv-contrib-python"]
+
+[project.urls]
+Repository = "https://github.com/city96/ComfyUI_ColorMod"
+
+[tool.comfy]
+PublisherId = "city96"
+DisplayName = "ComfyUI ColorMod"
+Icon = ""