HDR support

2024-04-03 17:38:40 +02:00
parent df5d1ca4f0
commit e1818dd45e
3 changed files with 196 additions and 0 deletions
@@ -21,6 +21,19 @@ else:
 		from .nodes.save_png import NODE_CLASS_MAPPINGS as save_png_nodes
 		NODE_CLASS_MAPPINGS.update(save_png_nodes)

+	# cv2 dep
+	try:
+		import cv2
+	except ImportError:
+		print("ColorMod: Can't find opencv! Please install to enable HDR/tonemapping support.")
+	else:
+		# HDR creation/etc nodes
+		from .nodes.hdr import NODE_CLASS_MAPPINGS as hdr_nodes
+		NODE_CLASS_MAPPINGS.update(hdr_nodes)
+		
+		# HDR save/load nodes
+		from .nodes.save_hdr import NODE_CLASS_MAPPINGS as save_hdr_nodes
+		NODE_CLASS_MAPPINGS.update(save_hdr_nodes)

 	# export
 	NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}
@@ -0,0 +1,97 @@
+import cv2
+import torch
+import numpy as np
+
+class HDRExposureFusion:
+	def __init__(self):
+		pass
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"image_a": ("IMAGE",),
+			},
+			"optional": {
+				"image_b": ("IMAGE",),
+				"image_c": ("IMAGE",),
+				"image_d": ("IMAGE",),
+			}
+		}
+
+	RETURN_TYPES = ("IMAGE",)
+	FUNCTION = "create_hdr"
+	CATEGORY = "ColorMod/hdr"
+	TITLE = "Exposure Fusion"
+
+	def create_hdr(self, image_a, image_b=None, image_c=None, image_d=None):
+		def img_to_cv2(img):
+			img = img.cpu().numpy()
+			img = img[:, :, ::-1] # PIL RGB to OpenCV BGR
+			img = (img * 255.0).astype(np.uint8)
+			return img
+
+		images = [x.clone() for x in [image_a, image_b, image_c, image_d] if x is not None]
+		assert all([x.shape[0] == images[0].shape[0] for x in images[1:]]), "Batch size mismatch!"
+		images = torch.stack(images, dim=1)
+
+		out = []
+		for batch in images:
+			batch = [img_to_cv2(x) for x in batch]
+			hdr = cv2.createMergeMertens().process(batch)
+			out.append(
+				torch.from_numpy(hdr[:, :, ::-1].copy())
+			)
+		out = torch.stack(out, dim=0)
+		return (out,)
+
+class HDRCreate:
+	def __init__(self):
+		pass
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"image_a": ("IMAGE",),
+				"image_b": ("IMAGE",),
+				"image_c": ("IMAGE",),
+				"exposure_a": ("FLOAT", {"default": 1.0, "min":  0.001, "max": 1024.0, "step": 0.1}),
+				"exposure_b": ("FLOAT", {"default": 2.5, "min":  0.001, "max": 1024.0, "step": 0.1}),
+				"exposure_c": ("FLOAT", {"default": 8.0, "min":  0.001, "max": 1024.0, "step": 0.1}),
+			}
+		}
+
+	RETURN_TYPES = ("IMAGE",)
+	FUNCTION = "create_hdr"
+	CATEGORY = "ColorMod/hdr"
+	TITLE = "Create HDR image"
+
+	def create_hdr(self, image_a, image_b, image_c, exposure_a, exposure_b, exposure_c):
+		def img_to_cv2(img):
+			img = img.cpu().numpy()
+			img = img[:, :, ::-1] # PIL RGB to OpenCV BGR
+			img = (img * 255.0).astype(np.uint8)
+			return img
+
+		images = [x.clone() for x in [image_a, image_b, image_c]]
+		assert all([x.shape[0] == images[0].shape[0] for x in images[1:]]), "Batch size mismatch!"
+		images = torch.stack(images, dim=1)
+		times = [exposure_a, exposure_b, exposure_c]
+		times = np.array(times, dtype=np.float32)
+		
+		out = []
+		for batch in images:
+			batch = [img_to_cv2(x) for x in batch]
+			cal = cv2.createCalibrateDebevec().process(batch, times)
+			hdr = cv2.createMergeDebevec().process(batch, times, cal)
+			out.append(
+				torch.from_numpy(hdr[:, :, ::-1].copy())
+			)
+		out = torch.stack(out, dim=0)
+		return (out,)
+
+NODE_CLASS_MAPPINGS = {
+	"HDRExposureFusion": HDRExposureFusion,
+	"HDRCreate": HDRCreate,
+}
@@ -0,0 +1,86 @@
+import os
+import cv2
+import torch
+import hashlib
+import numpy as np
+
+import folder_paths
+
+class SaveImageHDR:
+	def __init__(self):
+		self.output_dir = folder_paths.get_output_directory()
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"images": ("IMAGE", ),
+				"filename_prefix": ("STRING", {"default": "HDR/ComfyUI"})
+			}
+		}
+
+	OUTPUT_NODE = True
+	RETURN_TYPES = ()
+	FUNCTION = "save_image"
+	CATEGORY = "ColorMod/hdr"
+	TITLE = "Save Image (HDR)"
+
+	def save_image(self, images, filename_prefix):
+		full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
+			filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
+		)
+		for image in images:
+			file = f"{filename}_{counter:05}_.hdr"
+			path = os.path.join(full_output_folder, file)
+			image = image.cpu().numpy()
+			image = image[:, :, ::-1]
+			cv2.imwrite(path, image)
+		return ()
+
+class LoadImageHDR:
+	def __init__(self):
+		pass
+
+	@classmethod
+	def INPUT_TYPES(s):
+		exts = [".hdr"]
+		input_dir = folder_paths.get_input_directory()
+		files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+		files = [f for f in files if any([f.endswith(x) for x in exts])]
+		return {
+			"required" : {
+				"image": (sorted(files),)
+			}
+		}
+
+	RETURN_TYPES = ("IMAGE",)
+	FUNCTION = "load_image"
+	CATEGORY = "ColorMod/hdr"
+	TITLE = "Load Image (HDR)"
+
+	def load_image(self, image):
+		path = folder_paths.get_annotated_filepath(image)
+		img = cv2.imread(path, -1)
+		assert img is not None, "Failed to read image!"
+		out = torch.from_numpy(img[:, :, ::-1].copy()).unsqueeze(0)
+		print(f"Loaded HDR image [{torch.min(out)},{torch.max(out)}]")
+		return (out,)
+
+	@classmethod
+	def IS_CHANGED(s, image):
+		image_path = folder_paths.get_annotated_filepath(image)
+		m = hashlib.sha256()
+		with open(image_path, 'rb') as f:
+			m.update(f.read())
+		return m.digest().hex()
+
+	@classmethod
+	def VALIDATE_INPUTS(s, image):
+		if not folder_paths.exists_annotated_filepath(image):
+			return "Invalid image file: {}".format(image)
+		return True
+
+NODE_CLASS_MAPPINGS = {
+	"SaveImageHDR": SaveImageHDR,
+	"LoadImageHDR": LoadImageHDR,
+}