diff --git a/Exposure2HDR.py b/Exposure2HDR.py index 1221ac7..6a2a805 100644 --- a/Exposure2HDR.py +++ b/Exposure2HDR.py @@ -3,7 +3,7 @@ import torch class Exposure2HDR: """ - DiffusionLight Ball2Envmap class + DiffusionLight Exposure2HDR class """ def __init__(self): diff --git a/README.md b/README.md index 299f6e8..0b2f999 100644 --- a/README.md +++ b/README.md @@ -4,5 +4,5 @@ DiffusionLight node implmented to ComfyUI ## PLAN - [X] Ball2Envmap - [X] Exposure2HDR -- [ ] save HDR image +- [X] save HDR image - [ ] view HDR image using somekind of slider or generate list of differnet EVs \ No newline at end of file diff --git a/SaveHDR.py b/SaveHDR.py new file mode 100644 index 0000000..b8b6e23 --- /dev/null +++ b/SaveHDR.py @@ -0,0 +1,43 @@ +import numpy as np +import torch +import cv2 +import imageio +import folder_paths + +class SaveHDR: + """ + DiffusionLight SaveHDR class + + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "hdr_image": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "DiffusionLight"}), + "file_extension": (["hdr","npy", "exr"], {"default": "hdr"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "save_hdr" + + def save_hdr(self, hdr_image, filename_prefix, file_extension): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory()) + filename = f"{filename_prefix}_{counter:04d}.{file_extension}" + full_path = f"{full_output_folder}/{filename}" + print(f"Saving HDR image to {full_path}") + out_image = hdr_image[0].cpu().numpy().astype(np.float32) + if file_extension == "npy": + np.save(full_path, out_image) + elif file_extension == "hdr": + imageio.imwrite(full_path, out_image) + else: + cv2.imwrite(full_path, out_image) # save in HDR format + return (hdr_image, ) + + diff --git a/__init__.py b/__init__.py index b17a8e1..66b39dc 100644 --- a/__init__.py +++ b/__init__.py @@ -1,16 +1,19 @@ from .Ball2Envmap import Ball2Envmap from .Exposure2HDR import Exposure2HDR +from .SaveHDR import SaveHDR # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "DiffusionLightBall2Envmap": Ball2Envmap, - "DiffusionLightExposure2HDR": Exposure2HDR + "DiffusionLightExposure2HDR": Exposure2HDR, + "DiffusionLightSaveHDR": SaveHDR } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "DiffusionLightBall2Envmap": "Ball2Envmap", - "DiffusionLightExposure2HDR": "Exposure2HDR" + "DiffusionLightExposure2HDR": "Exposure2HDR", + "DiffusionLightSaveHDR": "SaveHDR" }