nodes/aimerib-ComfyUI_HigherBitDepthSaveImage
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update bit-depth conversion logic

Browse Source 
add initial icc profile support
This commit is contained in:
Aimeri Baddouh
2024-09-13 09:54:33 -04:00
parent c899e4b8c8
commit ef1683bb65
2 changed files with 144 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+7 -4
View File
@@ -6,11 +6,14 @@ This is a ComfyUI node that allows you to save images with a higher bit depth.
2. Place the file in the `custom_nodes` directory of your ComfyUI installation.
## Usage
1. Add the `Save Image (Higher Bit Depth)` node to your graph.
2. Connect the image you want to save to the `IMAGE` input of the node.
3. Set the `filename_prefix` to the directory where you want to save the image.
1. Add the `Save Image Higher Bit Depth` node to your graph.
2. Connect the image you want to save to the `IMAGE` input of the node. (Can be used in combination with other nodes that output images, like `SaveImage`)
3. Set the `filename_prefix` to the desired prefix for the filename of the image you want to save.
4. Set the `bit_depth` to the bit depth of the image you want to save.
5. Set the `format` to the format of the image you want to save. Default is `PNG`.
4.1 For 8bits (8-bit per channel - 24bit RGB) images, set `bit_depth` to `8`.
4.2 For 16bits (16-bit per channel - 48bit RGB) images, set `bit_depth` to `16`.
4.3 For 32bits (32-bit per channel - 96bit RGB) images, set `bit_depth` to `32`.
5. Set the `format` to the format of the image you want to save. Can be `PNG`, `TIFF`, or `JPEG`. Default is `PNG`.
save_image_node.py
+137 -76
View File
@@ -1,4 +1,4 @@
from PIL import Image
from PIL import Image, ImageCms
import numpy as np
import cv2
import folder_paths
@@ -18,13 +18,19 @@ class SaveImageHigherBitDepth:
@classmethod
def INPUT_TYPES(s):
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))]
return {
"required": {
"images": ("IMAGE", {"tooltip": "The images to save."}),
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix of the filename of the image you want to save."}),
"bit_depth": (["8", "16", "32"], {"default": "8", "tooltip": "The bit depth of the image you want to save."}),
"format": (["PNG", "TIFF", "JPEG"], {"default": "PNG", "tooltip": "The format of the image you want to save."})
"format": (["PNG", "TIFF", "JPEG"], {"default": "PNG", "tooltip": "The format of the image you want to save."}),
},
"optional": {
"new_icc_profile": (["LAB", "XYZ", "sRGB", "Default"], {"default": "Default", "tooltip": "The ICC profile to embed in the saved image.\n'LAB', 'XYZ', and 'sRGB' are color spaces.\n'LAB' is a color space that is designed to be perceptually uniform, meaning that the color difference between two points in the color space is proportional to the actual color difference perceived by the human eye.\n'XYZ' is a color space that is used in the CIE standard for colorimetry.\n'sRGB' is a color space that is widely used in digital imaging and is supported by most web browsers and operating systems."}),
"existing_icc_profile_path": ("STRING", {"default": "", "tooltip": "The custom ICC profile to embed in the saved image.\nThe path to the ICC profile file."})
}
}
RETURN_TYPES = ()
@@ -35,84 +41,139 @@ class SaveImageHigherBitDepth:
CATEGORY = "image"
DESCRIPTION = "Save image with higher bit depth and format."
def save_images(self, images, bit_depth, format="PNG", filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
filename_prefix += self.prefix_append
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])
results = list()
for (batch_number, image) in enumerate(images):
img_np = image.cpu().numpy()
if img_np.ndim != 3 or img_np.shape[2] != 3:
raise ValueError("Expected an RGB image with 3 channels.")
if bit_depth == "32":
format = "TIFF"
img = img_np.astype(np.float32)
# img = self.convert_opencv_to_pil(img_np)
elif bit_depth == "16":
format = "TIFF" if format == "TIFF" else "PNG"
img = np.clip(img_np * 65535.0, 0, 65535).astype(np.uint16)
# img = self.convert_opencv_to_pil(img_np)
else:
i = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
img = Image.fromarray(i, mode='RGB')
metadata = None
if not args.disable_metadata and format == "PNG":
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
# file = f"{filename_with_batch_num}_{counter:05}_.png"
# img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
if format == "JPEG":
file = f"{filename_with_batch_num}_{counter:05}_.jpg"
img.save(os.path.join(full_output_folder, file))
elif format == "PNG" and bit_depth == "8":
file = f"{filename_with_batch_num}_{counter:05}_.png"
img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
elif format == "PNG":
file = f"{filename_with_batch_num}_{counter:05}_.png"
cv2.imwrite(os.path.join(full_output_folder, file), img)
# img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
elif format == "TIFF":
file = f"{filename_with_batch_num}_{counter:05}_.tiff"
cv2.imwrite(os.path.join(full_output_folder, file), img)
# img.save(os.path.join(full_output_folder, file))
results.append({
"filename": file,
"subfolder": subfolder,
"type": self.type
})
counter += 1
return { "ui": { "images": results } }
def convert_opencv_to_pil(self, img_np):
def save_images(self, images, bit_depth, new_icc_profile="Default", existing_icc_profile_path="", format="PNG", filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
try:
if img_np.dtype == np.float32:
# For 32-bit float images
img_np = (img_np * 65535).astype(np.uint16)
elif img_np.dtype == np.uint8:
# For 8-bit images, no change needed
pass
elif img_np.dtype == np.uint16:
# For 16-bit images, no change needed
pass
else:
raise ValueError(f"Unsupported image dtype: {img_np.dtype}")
if new_icc_profile != "Default" and existing_icc_profile_path != "":
raise ValueError("Cannot use both new_icc_profile and existing_icc_profile_path. Please provide only one.")
# Convert BGR to RGB
img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
filename_prefix += self.prefix_append
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])
results = list()
# Create PIL Image directly from numpy array
img = Image.fromarray(img_np)
for (batch_number, image) in enumerate(images):
print(f"Processing image {batch_number + 1} of {len(images)}")
return img
img_np = image.cpu().numpy()
print(f"Image shape: {img_np.shape}, dtype: {img_np.dtype}")
print(f"Image min: {np.min(img_np)}, max: {np.max(img_np)}")
log_channel_stats(img_np, "Input")
save_debug_image(img_np, "input", batch_number, full_output_folder)
if img_np.ndim != 3 or img_np.shape[2] != 3:
raise ValueError(f"Expected an RGB image with 3 channels, got shape {img_np.shape}")
# Normalize input if necessary
if img_np.dtype == np.float32 or img_np.dtype == np.float64:
if np.min(img_np) < 0 or np.max(img_np) > 1:
print("Input values out of expected range [0, 1], normalizing")
img_np = (img_np - np.min(img_np)) / (np.max(img_np) - np.min(img_np))
log_channel_stats(img_np, "After normalization")
save_debug_image(img_np, "normalized", batch_number, full_output_folder)
# Convert to appropriate bit depth
if bit_depth == "32":
format = "TIFF"
# img_np = (img_np * 65535).astype(np.uint16)
if img_np.max() > 1: # Check if the input is in [0, 255] range
print("Input detected in [0, 255] range. Normalizing to [0, 1].")
img_np = img_np / 255.0
else:
print("Input detected in [0, 1] range. No normalization needed.")
img_np = img_np.astype(np.float32)
elif bit_depth == "16":
format = "TIFF" if format == "TIFF" else "PNG"
if img_np.max() > 1:
print("Input detected in [0, 255] range. Normalizing to [0, 1].")
img_np = img_np / 255.0
else:
print("Input detected in [0, 1] range. No normalization needed.")
img_np = np.clip(img_np * 65535.0, 0, 65535).astype(np.uint16)
else: # 8-bit
if img_np.dtype == np.uint8 and img_np.max() <= 255 and img_np.min() >= 0:
print("Input is already in 8-bit format. No conversion needed.")
elif img_np.max() <= 1.0 and img_np.min() >= 0:
print("Input detected in [0, 1] float range. Converting to [0, 255] uint8.")
img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
else:
print("Input in non-standard range. Normalizing to [0, 255] uint8.")
img_np = np.clip((img_np - img_np.min()) / (img_np.max() - img_np.min()) * 255.0, 0, 255).astype(np.uint8)
# img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
print(f"Converted image shape: {img_np.shape}, dtype: {img_np.dtype}")
print(f"Converted image min: {np.min(img_np)}, max: {np.max(img_np)}")
log_channel_stats(img_np, "After bit depth conversion")
save_debug_image(img_np, "bit_depth_converted", batch_number, full_output_folder)
if img_np[:,:,0].mean() < img_np[:,:,2].mean():
print("Detected BGR format, converting to RGB")
return cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img_np, mode='RGB')
# Handle ICC profile
if new_icc_profile != "Default":
icc_profile_cms_obj = ImageCms.createProfile(new_icc_profile)
icc_profile_path = ImageCms.getProfileName(icc_profile_cms_obj)
icc_profile = ImageCms.getOpenProfile(icc_profile_path)
img.info['icc_profile'] = icc_profile.tobytes()
elif existing_icc_profile_path:
icc_profile = ImageCms.getOpenProfile(existing_icc_profile_path)
img.info['icc_profile'] = icc_profile.tobytes()
# Save the image
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
log_channel_stats(img_np, "Before final save")
save_debug_image(img_np, "final", batch_number, full_output_folder)
file = self.save_image(img, format, bit_depth, full_output_folder, filename_with_batch_num, counter, prompt, extra_pnginfo)
results.append({
"filename": file,
"subfolder": subfolder,
"type": self.type
})
counter += 1
return { "ui": { "images": results } }
except Exception as e:
print(f"Error converting OpenCV image to PIL: {e}")
print(f"An error occurred while saving images: {str(e)}")
raise
def save_image(self, img, format, bit_depth, full_output_folder, filename_with_batch_num, counter, metadata=None, icc_profile=None):
extensions = {"JPEG": "jpg", "PNG": "png", "TIFF": "tiff"}
file = f"{filename_with_batch_num}_{counter:05}_.{extensions[format]}"
full_path = os.path.join(full_output_folder, file)
save_kwargs = {}
if icc_profile:
save_kwargs['icc_profile'] = img.info.get('icc_profile')
if format == "PNG":
if bit_depth == "8":
img.save(full_path, format="PNG", pnginfo=metadata, compress_level=self.compress_level, **save_kwargs)
else:
img.save(full_path, format="PNG", **save_kwargs)
elif format == "JPEG":
img.save(full_path, format="JPEG", **save_kwargs)
elif format == "TIFF":
img.save(full_path, format="TIFF", **save_kwargs)
return file
def save_debug_image(img_np, stage, batch_number, full_output_folder):
debug_img = Image.fromarray((img_np * 255).astype(np.uint8) if img_np.dtype != np.uint8 else img_np)
debug_filename = f"debug_{stage}_batch{batch_number}.png"
debug_path = os.path.join(full_output_folder, debug_filename)
debug_img.save(debug_path)
print(f"Saved debug image for {stage}: {debug_path}")
def log_channel_stats(img_np, stage):
for i, channel in enumerate(['R', 'G', 'B']):
channel_data = img_np[:,:,i]
print(f"{stage} - {channel} channel: min={np.min(channel_data):.4f}, max={np.max(channel_data):.4f}, mean={np.mean(channel_data):.4f}")