180 lines
9.2 KiB
Python
180 lines
9.2 KiB
Python
from PIL import Image, ImageCms
|
|
import numpy as np
|
|
import cv2
|
|
import folder_paths
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import os
|
|
import json
|
|
|
|
from comfy.cli_args import args
|
|
import comfy.utils
|
|
|
|
class SaveImageHigherBitDepth:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
self.compress_level = 4
|
|
|
|
@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."}),
|
|
},
|
|
"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 = ()
|
|
FUNCTION = "save_images"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "image"
|
|
DESCRIPTION = "Save image with higher bit depth and format."
|
|
|
|
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 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.")
|
|
|
|
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):
|
|
print(f"Processing image {batch_number + 1} of {len(images)}")
|
|
|
|
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"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}")
|