diff --git a/__init__.py b/__init__.py index a4f9427..be69b4b 100644 --- a/__init__.py +++ b/__init__.py @@ -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()} diff --git a/nodes/hdr.py b/nodes/hdr.py new file mode 100644 index 0000000..c46f656 --- /dev/null +++ b/nodes/hdr.py @@ -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, +} diff --git a/nodes/save_hdr.py b/nodes/save_hdr.py new file mode 100644 index 0000000..592e9d6 --- /dev/null +++ b/nodes/save_hdr.py @@ -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, +}