Add colorspace conversion nodes

__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"]

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,
+}

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]

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 = ""