Finish most formats. Missing 16bit png still.
add text fixtures
This commit is contained in:
Aimeri Baddouh
@@ -0,0 +1,25 @@
|
||||
# folder_paths.py
|
||||
# Utility functions to test the output of the save_image_node.py
|
||||
|
||||
import os
|
||||
|
||||
def get_output_directory():
|
||||
return os.path.abspath("output")
|
||||
|
||||
def get_save_image_path(filename_prefix, output_dir, width, height):
|
||||
# Create output directory if it doesn't exist
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Define filename pattern
|
||||
filename = f"{filename_prefix}_%batch_num%_image"
|
||||
|
||||
# Initialize counter
|
||||
counter = 0
|
||||
|
||||
# Define subfolder (optional)
|
||||
subfolder = ""
|
||||
|
||||
return output_dir, filename, counter, subfolder, filename_prefix
|
||||
|
||||
def get_input_directory():
|
||||
return os.path.abspath("input")
|
||||
@@ -0,0 +1,5 @@
|
||||
tifffile
|
||||
imageio
|
||||
pypng
|
||||
piexif
|
||||
opencv-python-headless
|
||||
+408
-56
@@ -5,16 +5,184 @@ import folder_paths
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
import os
|
||||
import json
|
||||
import tifffile
|
||||
import png # pypng
|
||||
import imageio
|
||||
import zlib
|
||||
import piexif # For JPEG EXIF metadata
|
||||
|
||||
from comfy.cli_args import args
|
||||
import comfy.utils
|
||||
# from comfy.cli_args import args
|
||||
# import comfy.utils
|
||||
|
||||
def get_profile_description(icc_profile):
|
||||
"""
|
||||
Extracts the ICC profile description.
|
||||
|
||||
Args:
|
||||
icc_profile (ImageCmsProfile): The opened ICC profile.
|
||||
|
||||
Returns:
|
||||
str: The profile description or 'sRGB' if extraction fails.
|
||||
"""
|
||||
try:
|
||||
# Access the 'description' attribute from the profile
|
||||
description = icc_profile.profile.profile_description
|
||||
# Ensure the profile name is within PNG's limitations (1-79 bytes)
|
||||
if len(description.encode('latin1')) > 79:
|
||||
print("Profile description too long for PNG iCCP chunk. Defaulting to 'sRGB'.")
|
||||
return 'sRGB'
|
||||
return description
|
||||
except Exception as e:
|
||||
print(f"Could not extract profile description: {e}. Defaulting to 'sRGB'.")
|
||||
return 'sRGB'
|
||||
|
||||
def create_iccp_chunk(icc_profile_bytes, profile_name='sRGB'):
|
||||
"""
|
||||
Creates the iCCP chunk data required for embedding ICC profiles in PNGs.
|
||||
|
||||
Args:
|
||||
icc_profile_bytes (bytes): Raw ICC profile data.
|
||||
profile_name (str): Name of the ICC profile. Must be 1-79 bytes.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the chunk type and data.
|
||||
"""
|
||||
if len(profile_name.encode('latin1')) > 79:
|
||||
raise ValueError("Profile name too long for iCCP chunk (max 79 bytes).")
|
||||
|
||||
# Encode profile name and ensure it's null-terminated
|
||||
profile_name_encoded = profile_name.encode('latin1') + b'\x00'
|
||||
|
||||
# Compression method is 0 (deflate)
|
||||
compression_method = b'\x00'
|
||||
|
||||
# Compress the ICC profile data using zlib
|
||||
compressed_icc = zlib.compress(icc_profile_bytes)
|
||||
|
||||
# Combine to form the iCCP chunk data
|
||||
iccp_data = profile_name_encoded + compression_method + compressed_icc
|
||||
|
||||
return ('iCCP', iccp_data)
|
||||
|
||||
def save_png_with_icc(full_path, img_np, icc_profile_bytes, profile_name='sRGB', metadata_dict=None):
|
||||
"""
|
||||
Saves a 16-bit PNG with an embedded ICC profile and additional metadata using pypng.
|
||||
|
||||
Args:
|
||||
full_path (str): Path to save the PNG file.
|
||||
img_np (numpy.ndarray): Image data as a NumPy array with dtype uint16.
|
||||
icc_profile_bytes (bytes): Raw ICC profile data.
|
||||
profile_name (str): Name of the ICC profile.
|
||||
metadata_dict (dict, optional): Additional metadata to embed.
|
||||
|
||||
Raises:
|
||||
ValueError: If the image data is not in the correct format.
|
||||
"""
|
||||
if img_np.dtype != np.uint16:
|
||||
print(f"Image data dtype: {img_np.dtype}")
|
||||
raise ValueError("Image data must be uint16 for 16-bit PNGs.")
|
||||
if img_np.ndim != 3 or img_np.shape[2] != 3:
|
||||
raise ValueError("Image data must have 3 channels (RGB).")
|
||||
|
||||
height, width, channels = img_np.shape
|
||||
# Reshape the image to a list of rows
|
||||
rows = img_np.reshape(-1, channels).tolist()
|
||||
|
||||
# Create the iCCP chunk
|
||||
chunks = []
|
||||
if icc_profile_bytes:
|
||||
iccp_chunk = create_iccp_chunk(icc_profile_bytes, profile_name)
|
||||
chunks.append(iccp_chunk)
|
||||
|
||||
# Add additional metadata as text chunks
|
||||
if metadata_dict:
|
||||
for key, value in metadata_dict.items():
|
||||
# Each key-value pair is added as a separate text chunk
|
||||
text_chunk = ('tEXt', f"{key}\x00{value}".encode('latin1'))
|
||||
chunks.append(text_chunk)
|
||||
|
||||
# Initialize the PNG writer
|
||||
writer = png.Writer(width=width, height=height, bitdepth=16, compression=9, greyscale=False)
|
||||
|
||||
# Write the PNG file with the iCCP and additional metadata chunks
|
||||
with open(full_path, 'wb') as f:
|
||||
writer.write(f, rows, chunks=chunks)
|
||||
|
||||
print(f"Saved 16-bit PNG with ICC profile and metadata at: {full_path}")
|
||||
|
||||
|
||||
def embed_metadata_tiff(full_path, img_np, metadata_dict):
|
||||
"""
|
||||
Embeds metadata into a TIFF file using tifffile.
|
||||
|
||||
Args:
|
||||
full_path (str): Path to save the TIFF file.
|
||||
img_np (numpy.ndarray): Image data as a NumPy array.
|
||||
metadata_dict (dict): Dictionary of metadata to embed.
|
||||
|
||||
Raises:
|
||||
ValueError: If metadata_dict is not a dictionary.
|
||||
"""
|
||||
if not isinstance(metadata_dict, dict):
|
||||
raise ValueError("Metadata must be a dictionary.")
|
||||
|
||||
# Convert metadata_dict to JSON string
|
||||
metadata_json = json.dumps(metadata_dict)
|
||||
|
||||
# Prepare the metadata
|
||||
metadata = {'Software': 'ComfyUI', 'Description': metadata_json}
|
||||
|
||||
# Save using tifffile with metadata
|
||||
tifffile.imwrite(
|
||||
full_path,
|
||||
img_np,
|
||||
photometric='rgb',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
print(f"Saved TIFF with embedded metadata at: {full_path}")
|
||||
|
||||
def embed_metadata_jpeg(full_path, img_pil, metadata_dict):
|
||||
"""
|
||||
Embeds metadata into a JPEG file using piexif.
|
||||
|
||||
Args:
|
||||
full_path (str): Path to save the JPEG file.
|
||||
img_pil (PIL.Image.Image): PIL Image object.
|
||||
metadata_dict (dict): Dictionary of metadata to embed.
|
||||
|
||||
Raises:
|
||||
ValueError: If metadata_dict is not a dictionary.
|
||||
"""
|
||||
if not isinstance(metadata_dict, dict):
|
||||
raise ValueError("Metadata must be a dictionary.")
|
||||
|
||||
# Convert metadata_dict to JSON string
|
||||
metadata_json = json.dumps(metadata_dict)
|
||||
|
||||
# Create EXIF data
|
||||
exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "thumbnail": None}
|
||||
|
||||
# Embed the JSON metadata in the UserComment tag (Exif tag 37510)
|
||||
# UserComment requires a specific format: 8-byte prefix + comment
|
||||
prefix = b"ComfyUI\x00\x00\x00" # Custom prefix to identify the source
|
||||
user_comment = prefix + metadata_json.encode('utf-8')
|
||||
exif_dict["Exif"][piexif.ExifIFD.UserComment] = user_comment
|
||||
|
||||
# Convert exif_dict to bytes
|
||||
exif_bytes = piexif.dump(exif_dict)
|
||||
|
||||
# Save the JPEG with EXIF data
|
||||
img_pil.save(full_path, "JPEG", exif=exif_bytes, quality=95)
|
||||
|
||||
print(f"Saved JPEG with embedded metadata at: {full_path}")
|
||||
|
||||
class SaveImageHigherBitDepth:
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
self.compress_level = 4
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
self.compress_level = 4
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -30,7 +198,10 @@ class SaveImageHigherBitDepth:
|
||||
"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."})
|
||||
}
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
@@ -47,7 +218,20 @@ class SaveImageHigherBitDepth:
|
||||
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])
|
||||
# Adjust the order of dimensions based on assumption
|
||||
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]
|
||||
)
|
||||
# if img_shape[0] == 3:
|
||||
# # Likely (channels, height, width)
|
||||
# elif img_shape[2] == 3:
|
||||
# # Already (height, width, channels)
|
||||
# full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
# filename_prefix, self.output_dir, img_shape[1], img_shape[0]
|
||||
# )
|
||||
# else:
|
||||
# raise ValueError(f"Unexpected image shape: {img_shape}")
|
||||
|
||||
results = list()
|
||||
|
||||
for (batch_number, image) in enumerate(images):
|
||||
@@ -55,30 +239,55 @@ class SaveImageHigherBitDepth:
|
||||
|
||||
img_np = image.cpu().numpy()
|
||||
|
||||
# Detect and transpose if necessary
|
||||
if img_np.shape[0] == 3:
|
||||
img_np = np.transpose(img_np, (1, 2, 0)) # Transpose to (height, width, channels)
|
||||
print("Transposed image from (channels, height, width) to (height, width, channels).")
|
||||
elif img_np.shape[2] == 3:
|
||||
pass # Already in (height, width, channels)
|
||||
else:
|
||||
raise ValueError(f"Unexpected image shape after conversion: {img_np.shape}")
|
||||
|
||||
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)
|
||||
|
||||
save_debug_image(img_np, "input", counter, 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 img_np.dtype in [np.float32, 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)
|
||||
save_debug_image(img_np, "normalized", counter, full_output_folder)
|
||||
|
||||
# Initialize ICC profile bytes and profile name
|
||||
icc_profile_bytes = None
|
||||
profile_name = 'sRGB' # Default profile name
|
||||
|
||||
# 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_cms_obj)
|
||||
icc_profile_bytes = icc_profile.tobytes()
|
||||
profile_name = get_profile_description(icc_profile)
|
||||
elif existing_icc_profile_path:
|
||||
icc_profile = ImageCms.getOpenProfile(existing_icc_profile_path)
|
||||
icc_profile_bytes = icc_profile.tobytes()
|
||||
profile_name = get_profile_description(icc_profile)
|
||||
|
||||
# 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
|
||||
format = "TIFF" # TIFF supports higher bit depths
|
||||
# Ensure the data is in float32
|
||||
if img_np.max() > 1.0: # 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:
|
||||
@@ -87,13 +296,14 @@ class SaveImageHigherBitDepth:
|
||||
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
|
||||
if img_np.max() <= 1.0:
|
||||
print("Input detected in [0, 1] range. Scaling to [0, 65535].")
|
||||
img_np = img_np * 65535.0
|
||||
else:
|
||||
print("Input detected in [0, 1] range. No normalization needed.")
|
||||
print("Input detected in [0, 255] range. Scaling to [0, 65535].")
|
||||
img_np = img_np * (65535.0 / 255.0)
|
||||
|
||||
img_np = np.clip(img_np * 65535.0, 0, 65535).astype(np.uint16)
|
||||
img_np = np.clip(img_np, 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.")
|
||||
@@ -103,35 +313,44 @@ class SaveImageHigherBitDepth:
|
||||
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)
|
||||
save_debug_image(img_np, "bit_depth_converted", counter, 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)
|
||||
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
||||
|
||||
img = Image.fromarray(img_np, mode='RGB')
|
||||
# Log before saving
|
||||
log_channel_stats(img_np, "Before final save")
|
||||
save_debug_image(img_np, "final", counter, full_output_folder)
|
||||
|
||||
# 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()
|
||||
# Prepare metadata for embedding
|
||||
metadata = None
|
||||
# if not args.disable_metadata:
|
||||
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]))
|
||||
metadata_dict = {}
|
||||
if metadata:
|
||||
metadata_dict = extract_metadata_from_pnginfo(metadata)
|
||||
|
||||
# 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)
|
||||
file = self.save_image(
|
||||
img_np,
|
||||
format,
|
||||
bit_depth,
|
||||
full_output_folder,
|
||||
filename_with_batch_num,
|
||||
counter,
|
||||
metadata,
|
||||
icc_profile_bytes=icc_profile_bytes,
|
||||
profile_name=profile_name,
|
||||
metadata_dict=metadata_dict # Pass the metadata dictionary
|
||||
)
|
||||
results.append({
|
||||
"filename": file,
|
||||
"subfolder": subfolder,
|
||||
@@ -140,40 +359,173 @@ class SaveImageHigherBitDepth:
|
||||
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):
|
||||
def save_image(self, img_np, format, bit_depth, full_output_folder, filename_with_batch_num, counter, extra_pnginfo=None, icc_profile_bytes=None, profile_name='sRGB', metadata_dict=None):
|
||||
"""
|
||||
Saves the image in the specified format and bit depth, embedding ICC profiles and metadata as needed.
|
||||
|
||||
Args:
|
||||
img (PIL.Image.Image): The PIL Image to save.
|
||||
format (str): The image format ('PNG', 'TIFF', 'JPEG').
|
||||
bit_depth (str): The bit depth ('8', '16', '32').
|
||||
full_output_folder (str): The directory to save the image.
|
||||
filename_with_batch_num (str): The filename with batch number.
|
||||
counter (int): The image counter.
|
||||
prompt (str, optional): Additional prompt information.
|
||||
extra_pnginfo (PngInfo, optional): Extra PNG metadata.
|
||||
icc_profile_bytes (bytes, optional): ICC profile bytes.
|
||||
profile_name (str, optional): Name of the ICC profile.
|
||||
metadata_dict (dict, optional): Metadata to embed in all formats.
|
||||
|
||||
Returns:
|
||||
str: The saved filename.
|
||||
"""
|
||||
extensions = {"JPEG": "jpg", "PNG": "png", "TIFF": "tiff"}
|
||||
file = f"{filename_with_batch_num}_{counter:05}_.{extensions[format]}"
|
||||
file = f"{filename_with_batch_num}.{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')
|
||||
print(f"Saving file: {full_path}")
|
||||
print(f"Format: {format}, Bit Depth: {bit_depth}, ICC Profile: {'Yes' if icc_profile_bytes else 'No'}, Profile Name: {profile_name}")
|
||||
|
||||
if format == "PNG":
|
||||
if bit_depth == "8":
|
||||
img.save(full_path, format="PNG", pnginfo=metadata, compress_level=self.compress_level, **save_kwargs)
|
||||
img = Image.fromarray(img_np, mode='RGB')
|
||||
|
||||
if format == "TIFF":
|
||||
# Handle metadata embedding for TIFF
|
||||
if metadata_dict:
|
||||
embed_metadata_tiff(full_path, img_np, metadata_dict)
|
||||
else:
|
||||
img.save(full_path, format="PNG", **save_kwargs)
|
||||
# Save without additional metadata
|
||||
tifffile.imwrite(
|
||||
full_path,
|
||||
img_np,
|
||||
photometric='rgb',
|
||||
metadata=None,
|
||||
)
|
||||
print(f"Saved TIFF without additional metadata at: {full_path}")
|
||||
elif format == "PNG" and bit_depth == "16":
|
||||
# Use pypng to save 16-bit PNG with ICC Profile and metadata
|
||||
if img_np.dtype != np.uint16:
|
||||
raise ValueError("Image data must be uint16 for 16-bit PNGs.")
|
||||
if img_np.ndim != 3 or img_np.shape[2] != 3:
|
||||
raise ValueError("Image data must have 3 channels (RGB).")
|
||||
if icc_profile_bytes:
|
||||
# Create a PIL Image from the NumPy array
|
||||
img.save(
|
||||
full_path,
|
||||
format='PNG',
|
||||
pnginfo=extra_pnginfo,
|
||||
icc_profile=icc_profile_bytes,
|
||||
compress_level=self.compress_level # Adjust compression level as needed
|
||||
)
|
||||
else:
|
||||
# Save without ICC Profile and metadata
|
||||
img.save(full_path, format="PNG", pnginfo=extra_pnginfo, compress_level=self.compress_level)
|
||||
print(f"Saved 16-bit PNG without ICC profile: {full_path}")
|
||||
elif format == "PNG" and bit_depth == "8":
|
||||
# Save 8-bit PNG with extra_pnginfo
|
||||
img.save(full_path, format="PNG", pnginfo=extra_pnginfo, compress_level=self.compress_level)
|
||||
print(f"Saved 8-bit PNG at: {full_path}")
|
||||
elif format == "JPEG":
|
||||
img.save(full_path, format="JPEG", **save_kwargs)
|
||||
elif format == "TIFF":
|
||||
img.save(full_path, format="TIFF", **save_kwargs)
|
||||
# Handle metadata embedding for JPEG
|
||||
if metadata_dict:
|
||||
embed_metadata_jpeg(full_path, img, metadata_dict)
|
||||
else:
|
||||
img.save(full_path, format="JPEG", quality=95)
|
||||
print(f"Saved JPEG without additional metadata at: {full_path}")
|
||||
else:
|
||||
# Handle other formats and bit depths with PIL
|
||||
save_kwargs = {}
|
||||
if icc_profile_bytes:
|
||||
save_kwargs['icc_profile'] = icc_profile_bytes
|
||||
print("ICC Profile embedded using PIL.")
|
||||
|
||||
img.save(full_path, format=format, **save_kwargs)
|
||||
print(f"Saved {format} at: {full_path}")
|
||||
|
||||
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"
|
||||
def save_debug_image(img_np, stage, counter, full_output_folder):
|
||||
"""
|
||||
Saves a debug image for the specified stage.
|
||||
|
||||
Args:
|
||||
img_np (numpy.ndarray): Image data as a NumPy array.
|
||||
stage (str): The processing stage.
|
||||
counter (int): The image counter.
|
||||
full_output_folder (str): The directory to save the debug image.
|
||||
"""
|
||||
if img_np.dtype in [np.float32, np.float64]:
|
||||
if img_np.max() <= 1.0:
|
||||
debug_img = Image.fromarray((img_np * 255).astype(np.uint8))
|
||||
else:
|
||||
debug_img = Image.fromarray((np.clip(img_np / np.max(img_np), 0, 1) * 255).astype(np.uint8))
|
||||
elif img_np.dtype == np.uint16:
|
||||
debug_img = Image.fromarray((img_np / 257).astype(np.uint8)) # Scale down to 8-bit
|
||||
else:
|
||||
debug_img = Image.fromarray(img_np.astype(np.uint8))
|
||||
|
||||
debug_filename = f"debug_{stage}_batch{counter}.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):
|
||||
"""
|
||||
Logs statistics for each color channel.
|
||||
|
||||
Args:
|
||||
img_np (numpy.ndarray): Image data as a NumPy array.
|
||||
stage (str): The processing 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}")
|
||||
|
||||
def extract_metadata_from_pnginfo(pnginfo: PngInfo) -> dict:
|
||||
"""
|
||||
Extracts metadata key-value pairs from a PngInfo object.
|
||||
|
||||
Args:
|
||||
pnginfo (PngInfo): The PngInfo object containing metadata.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary of metadata key-value pairs.
|
||||
"""
|
||||
metadata = {}
|
||||
for chunk_type, data, _ in pnginfo.chunks:
|
||||
if chunk_type == b'tEXt':
|
||||
try:
|
||||
key, value = data.split(b'\x00', 1)
|
||||
key = key.decode('latin-1')
|
||||
value = value.decode('latin-1')
|
||||
metadata[key] = value
|
||||
except ValueError:
|
||||
print("Invalid tEXt chunk format.")
|
||||
elif chunk_type == b'zTXt':
|
||||
try:
|
||||
key, compressed_value = data.split(b'\x00', 1)
|
||||
key = key.decode('latin-1')
|
||||
value = zlib.decompress(compressed_value).decode('latin-1')
|
||||
metadata[key] = value
|
||||
except (ValueError, zlib.error):
|
||||
print("Invalid zTXt chunk format or compression.")
|
||||
elif chunk_type == b'iTXt':
|
||||
try:
|
||||
# iTXt format: key\0compression_flag\0language\0translated_keyword\0text
|
||||
parts = data.split(b'\x00', 4)
|
||||
if len(parts) >= 5:
|
||||
key, compression_flag, lang, tkey, text = parts
|
||||
key = key.decode('latin-1')
|
||||
if compression_flag == b'\x01':
|
||||
# Text is compressed
|
||||
value = zlib.decompress(text).decode('utf-8')
|
||||
else:
|
||||
value = text.decode('utf-8')
|
||||
metadata[key] = value
|
||||
except (ValueError, zlib.error):
|
||||
print("Invalid iTXt chunk format or compression.")
|
||||
return metadata
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
# test_save_image.py
|
||||
# Test the save_image_node.py script
|
||||
|
||||
import torch
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
from save_image_node import SaveImageHigherBitDepth
|
||||
|
||||
def main():
|
||||
# Instantiate the class
|
||||
test = SaveImageHigherBitDepth()
|
||||
|
||||
# Create a synthetic image tensor (RGB, 256x256) with values in [0,1]
|
||||
img_tensor = torch.rand(3, 256, 256)
|
||||
|
||||
# Create extra_pnginfo with metadata
|
||||
png_info = {}
|
||||
png_info['Author'] = 'TestUser'
|
||||
png_info['Description'] = 'Synthetic Image for Testing'
|
||||
|
||||
# Call the save_images method
|
||||
test.save_images(
|
||||
images=[img_tensor],
|
||||
bit_depth='16',
|
||||
format='PNG',
|
||||
filename_prefix='TestImage',
|
||||
extra_pnginfo=png_info,
|
||||
new_icc_profile='XYZ',
|
||||
# existing_icc_profile_path="/Users/aimeri/Downloads/Adobe ICC Profiles (end-user)/RGB/AdobeRGB1998.icc"
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user