Completely rewrote the saver node to address a few bugs and clean up the code
This commit is contained in:
Conor-Collins
+3
-3
@@ -29,9 +29,9 @@ NODE_CLASS_MAPPINGS.update({
|
||||
|
||||
# Add display names for better UI presentation
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
"ImageLoader": "Image Loader",
|
||||
"LoadExr": "Load EXR",
|
||||
"SaverNode": "Save Image",
|
||||
"ImageLoader": "coco Loader",
|
||||
"LoadExr": "coco Load EXR",
|
||||
"SaverNode": "coco Saver",
|
||||
"LoadExrLayerByName": "Load EXR Layer by Name",
|
||||
"CryptomatteLayer": "Cryptomatte Layer",
|
||||
"ColorspaceNode": "Colorspace",
|
||||
|
||||
+118
-45
@@ -3,8 +3,8 @@ import { app } from "../../../scripts/app.js";
|
||||
app.registerExtension({
|
||||
name: "Coco.Saver",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
// IMPORTANT: This must match the key in NODE_CLASS_MAPPINGS
|
||||
if (nodeData.name !== "SaverNode") {
|
||||
// CRITICAL: Must match NODE_CLASS_MAPPINGS key
|
||||
if (nodeData.name !== "SaverNode") {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -14,68 +14,141 @@ app.registerExtension({
|
||||
nodeType.prototype.onNodeCreated = function() {
|
||||
const me = onNodeCreated?.apply(this);
|
||||
|
||||
console.log("Saver node created, setting up dynamic widgets");
|
||||
// Format specifications matching Python
|
||||
const FORMAT_SPECS = {
|
||||
"exr": {
|
||||
depths: ["16", "32"], // EXR only supports half and full float
|
||||
showWidgets: ["exr_compression", "bit_depth", "save_as_grayscale"]
|
||||
},
|
||||
"png": {
|
||||
depths: ["8", "16"], // PNG only supports integer formats
|
||||
showWidgets: ["bit_depth", "save_as_grayscale"]
|
||||
},
|
||||
"jpg": {
|
||||
depths: ["8"],
|
||||
showWidgets: ["quality"]
|
||||
},
|
||||
"webp": {
|
||||
depths: ["8"],
|
||||
showWidgets: ["quality"]
|
||||
},
|
||||
"tiff": {
|
||||
depths: ["8", "16", "32"],
|
||||
showWidgets: ["bit_depth", "save_as_grayscale"]
|
||||
}
|
||||
};
|
||||
|
||||
// Get all widgets
|
||||
const fileType = this.widgets.find(w => w.name === "file_type");
|
||||
const bitDepth = this.widgets.find(w => w.name === "bit_depth");
|
||||
const exrCompression = this.widgets.find(w => w.name === "exr_compression");
|
||||
const quality = this.widgets.find(w => w.name === "quality");
|
||||
const saveAsGrayscale = this.widgets.find(w => w.name === "save_as_grayscale");
|
||||
const widgets = {
|
||||
fileType: this.widgets.find(w => w.name === "file_type"),
|
||||
bitDepth: this.widgets.find(w => w.name === "bit_depth"),
|
||||
exrCompression: this.widgets.find(w => w.name === "exr_compression"),
|
||||
quality: this.widgets.find(w => w.name === "quality"),
|
||||
saveAsGrayscale: this.widgets.find(w => w.name === "save_as_grayscale")
|
||||
};
|
||||
|
||||
if (!fileType) {
|
||||
if (!widgets.fileType) {
|
||||
console.warn("Saver: file_type widget not found");
|
||||
return me;
|
||||
}
|
||||
|
||||
// Define the callback that updates widget visibility
|
||||
const updateWidgetVisibility = () => {
|
||||
console.log("Saver: Updating widget visibility for type:", fileType.value);
|
||||
// Store original bit depth options
|
||||
const allBitDepthOptions = ["8", "16", "32"];
|
||||
|
||||
// Update bit depth options based on format
|
||||
const updateBitDepthOptions = (format) => {
|
||||
if (!widgets.bitDepth) return;
|
||||
|
||||
// Hide all format-specific widgets first
|
||||
if (bitDepth) bitDepth.hidden = true;
|
||||
if (exrCompression) exrCompression.hidden = true;
|
||||
if (quality) quality.hidden = true;
|
||||
if (saveAsGrayscale) saveAsGrayscale.hidden = true;
|
||||
|
||||
// Show relevant widgets based on file type
|
||||
switch (fileType.value) {
|
||||
case "exr":
|
||||
if (exrCompression) exrCompression.hidden = false;
|
||||
if (bitDepth) bitDepth.hidden = false;
|
||||
if (saveAsGrayscale) saveAsGrayscale.hidden = false;
|
||||
break;
|
||||
case "png":
|
||||
case "tiff":
|
||||
if (bitDepth) bitDepth.hidden = false;
|
||||
if (saveAsGrayscale) saveAsGrayscale.hidden = false;
|
||||
break;
|
||||
case "jpg":
|
||||
case "webp":
|
||||
if (quality) quality.hidden = false;
|
||||
break;
|
||||
const spec = FORMAT_SPECS[format];
|
||||
if (spec && spec.depths) {
|
||||
// Update widget options
|
||||
widgets.bitDepth.options.values = spec.depths;
|
||||
|
||||
// If current value is not valid, reset to first valid option
|
||||
if (!spec.depths.includes(widgets.bitDepth.value)) {
|
||||
widgets.bitDepth.value = spec.depths[0];
|
||||
|
||||
// Visual feedback - briefly highlight the change
|
||||
const originalColor = widgets.bitDepth.color;
|
||||
widgets.bitDepth.color = "#ff6b6b";
|
||||
setTimeout(() => {
|
||||
widgets.bitDepth.color = originalColor;
|
||||
app.canvas.setDirty(true);
|
||||
}, 300);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Update node size to accommodate the visible widgets
|
||||
// Update widget visibility and options
|
||||
const updateWidgets = () => {
|
||||
const format = widgets.fileType.value;
|
||||
const spec = FORMAT_SPECS[format];
|
||||
|
||||
console.log(`Saver: Updating for format: ${format}`);
|
||||
|
||||
// Hide all optional widgets
|
||||
Object.keys(widgets).forEach(key => {
|
||||
if (key !== "fileType" && widgets[key]) {
|
||||
widgets[key].hidden = true;
|
||||
}
|
||||
});
|
||||
|
||||
// Show relevant widgets
|
||||
if (spec && spec.showWidgets) {
|
||||
spec.showWidgets.forEach(widgetName => {
|
||||
const camelCaseName = widgetName.replace(/_([a-z])/g, (_, letter) => letter.toUpperCase());
|
||||
if (widgets[camelCaseName]) {
|
||||
widgets[camelCaseName].hidden = false;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Update bit depth options
|
||||
updateBitDepthOptions(format);
|
||||
|
||||
// Resize node
|
||||
requestAnimationFrame(() => {
|
||||
this.setSize([this.size[0], this.computeSize([this.size[0], this.size[1]])[1]]);
|
||||
const newHeight = this.computeSize([this.size[0], this.size[1]])[1];
|
||||
this.setSize([this.size[0], newHeight]);
|
||||
app.canvas.setDirty(true);
|
||||
});
|
||||
};
|
||||
|
||||
// Attach the callback to the file_type widget
|
||||
const originalCallback = fileType.callback;
|
||||
fileType.callback = function(...args) {
|
||||
if (originalCallback) {
|
||||
originalCallback.apply(this, args);
|
||||
// Override file type callback
|
||||
const originalFileTypeCallback = widgets.fileType.callback;
|
||||
widgets.fileType.callback = function(...args) {
|
||||
if (originalFileTypeCallback) {
|
||||
originalFileTypeCallback.apply(this, args);
|
||||
}
|
||||
updateWidgetVisibility();
|
||||
updateWidgets();
|
||||
};
|
||||
|
||||
// Initial update with a small delay to ensure widgets are ready
|
||||
setTimeout(updateWidgetVisibility, 1);
|
||||
// Override bit depth callback for validation
|
||||
if (widgets.bitDepth) {
|
||||
const originalBitDepthCallback = widgets.bitDepth.callback;
|
||||
widgets.bitDepth.callback = function(value) {
|
||||
// Validate bit depth for current format
|
||||
const currentFormat = widgets.fileType.value;
|
||||
const spec = FORMAT_SPECS[currentFormat];
|
||||
|
||||
if (spec && !spec.depths.includes(value)) {
|
||||
console.warn(`Invalid bit depth ${value} for format ${currentFormat}`);
|
||||
// Reset to first valid option
|
||||
this.value = spec.depths[0];
|
||||
app.canvas.setDirty(true);
|
||||
return;
|
||||
}
|
||||
|
||||
if (originalBitDepthCallback) {
|
||||
originalBitDepthCallback.apply(this, arguments);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Initial setup with delay
|
||||
setTimeout(updateWidgets, 10);
|
||||
|
||||
return me;
|
||||
};
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
+194
-241
@@ -3,32 +3,35 @@ import torch
|
||||
import numpy as np
|
||||
import tifffile
|
||||
import folder_paths
|
||||
from typing import Dict, Tuple
|
||||
from typing import Dict, Tuple, Optional
|
||||
import OpenImageIO as oiio
|
||||
from datetime import datetime
|
||||
import sys
|
||||
|
||||
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
|
||||
import cv2 as cv
|
||||
|
||||
class saver:
|
||||
"""Image saver node with dynamic widget behavior"""
|
||||
"""Optimized image saver node with consistent bit depth handling"""
|
||||
|
||||
# Format specifications
|
||||
FORMAT_SPECS = {
|
||||
"exr": {"depths": [16, 32], "opencv": False}, # EXR only supports half and full float
|
||||
"png": {"depths": [8, 16], "opencv": False}, # PNG only supports integer formats
|
||||
"jpg": {"depths": [8], "opencv": True},
|
||||
"webp": {"depths": [8], "opencv": True},
|
||||
"tiff": {"depths": [8, 16, 32], "opencv": False}
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
# Common widgets that are always visible at the top
|
||||
"images": ("IMAGE",),
|
||||
"file_path": ("STRING", {"default": "ComfyUI"}),
|
||||
"file_path": ("STRING", {"default": ""}),
|
||||
"filename": ("STRING", {"default": "ComfyUI"}),
|
||||
"use_versioning": ("BOOLEAN", {"default": False}),
|
||||
"version": ("INT", {"default": 1, "min": -1, "max": 999}),
|
||||
"start_frame": ("INT", {"default": 1001, "min": 0, "max": 99999999}),
|
||||
"frame_pad": ("INT", {"default": 4, "min": 1, "max": 8}),
|
||||
|
||||
# This widget controls the visibility of others
|
||||
"file_type": (["exr", "png", "jpg", "webp", "tiff"], {"default": "png"}),
|
||||
|
||||
# These widgets will be shown/hidden based on file_type
|
||||
"bit_depth": (["8", "16", "32"], {"default": "16"}),
|
||||
"exr_compression": (["none", "zip", "zips", "rle", "pxr24", "b44", "b44a", "dwaa", "dwab"],
|
||||
{"default": "zips"}),
|
||||
@@ -45,277 +48,227 @@ class saver:
|
||||
FUNCTION = "save_images"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "COCO Tools/Savers"
|
||||
|
||||
# This tells ComfyUI that file_type widget needs custom handling
|
||||
CUSTOM_WIDGETS = ["file_type"]
|
||||
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
|
||||
@staticmethod
|
||||
def is_grayscale(image: np.ndarray) -> bool:
|
||||
"""Check if an image is grayscale by either being single channel or having identical RGB channels."""
|
||||
def is_grayscale_fast(image: np.ndarray, sample_rate: float = 0.1) -> bool:
|
||||
"""Fast grayscale detection by sampling pixels"""
|
||||
if len(image.shape) == 2 or image.shape[-1] == 1:
|
||||
return True
|
||||
if image.shape[-1] == 3:
|
||||
return np.allclose(image[..., 0], image[..., 1]) and np.allclose(image[..., 1], image[..., 2])
|
||||
# Sample 10% of pixels for large images
|
||||
total_pixels = image.shape[0] * image.shape[1]
|
||||
if total_pixels > 1000000: # 1MP
|
||||
indices = np.random.choice(total_pixels, int(total_pixels * sample_rate), replace=False)
|
||||
flat_img = image.reshape(-1, 3)
|
||||
sampled = flat_img[indices]
|
||||
return np.allclose(sampled[:, 0], sampled[:, 1], rtol=0.01) and \
|
||||
np.allclose(sampled[:, 1], sampled[:, 2], rtol=0.01)
|
||||
else:
|
||||
return np.allclose(image[..., 0], image[..., 1], rtol=0.01) and \
|
||||
np.allclose(image[..., 1], image[..., 2], rtol=0.01)
|
||||
return False
|
||||
|
||||
def _validate_format_bitdepth(self, file_type: str, bit_depth: int) -> Tuple[str, int]:
|
||||
valid_combinations = {
|
||||
"exr": [16, 32], # OpenEXR supports half and full float
|
||||
"png": [8, 16, 32], # Now supporting 32-bit PNGs through OpenImageIO
|
||||
"jpg": [8],
|
||||
"jpeg": [8],
|
||||
"webp": [8],
|
||||
"tiff": [8, 16, 32]
|
||||
}
|
||||
|
||||
if bit_depth not in valid_combinations[file_type]:
|
||||
sys.stderr.write(f"Warning: {file_type} format only supports {valid_combinations[file_type]} bit depth. Adjusting to {valid_combinations[file_type][0]} bit.\n")
|
||||
bit_depth = valid_combinations[file_type][0]
|
||||
|
||||
return file_type, bit_depth
|
||||
def validate_bit_depth(self, file_type: str, bit_depth: int) -> int:
|
||||
"""Validate and adjust bit depth for format"""
|
||||
valid_depths = self.FORMAT_SPECS[file_type]["depths"]
|
||||
if bit_depth not in valid_depths:
|
||||
return valid_depths[0]
|
||||
return bit_depth
|
||||
|
||||
def increment_filename(self, filepath: str) -> str:
|
||||
base, ext = os.path.splitext(filepath)
|
||||
counter = 1
|
||||
new_filepath = f"{base}_{counter:05d}{ext}"
|
||||
while os.path.exists(new_filepath):
|
||||
counter += 1
|
||||
new_filepath = f"{base}_{counter:05d}{ext}"
|
||||
return new_filepath
|
||||
|
||||
def _convert_bit_depth(self, img: np.ndarray, bit_depth: int) -> np.ndarray:
|
||||
"""Convert image to specified bit depth."""
|
||||
if bit_depth == 8:
|
||||
return (np.clip(img, 0, 1) * 255).astype(np.uint8)
|
||||
elif bit_depth == 16:
|
||||
if isinstance(img.dtype, np.floating):
|
||||
# For floating point data (like EXR), preserve as float16
|
||||
return img.astype(np.float16)
|
||||
else:
|
||||
# For integer data, scale to 16-bit range
|
||||
return (np.clip(img, 0, 1) * 65535).astype(np.uint16)
|
||||
else: # 32-bit
|
||||
return img.astype(np.float32)
|
||||
|
||||
def _prepare_image_for_saving(self, img: np.ndarray, file_type: str, save_as_grayscale: bool = False) -> np.ndarray:
|
||||
"""Prepare image for saving by handling color space and channel conversion."""
|
||||
# Handle single channel images
|
||||
if len(img.shape) == 2:
|
||||
img = img[..., np.newaxis]
|
||||
def convert_to_grayscale(self, img: np.ndarray) -> np.ndarray:
|
||||
"""Proper grayscale conversion using luminance weights"""
|
||||
if len(img.shape) == 2 or img.shape[-1] == 1:
|
||||
return img if img.shape[-1] == 1 else img[..., np.newaxis]
|
||||
|
||||
# Convert to grayscale if requested or if image is already grayscale
|
||||
if save_as_grayscale or self.is_grayscale(img):
|
||||
if img.shape[-1] == 3:
|
||||
# For EXR workflow, just take R channel to preserve original values
|
||||
img = img[..., 0:1]
|
||||
# Convert to BGR for OpenCV formats only (jpg, jpeg, webp) if not saving as grayscale
|
||||
elif file_type in ['jpg', 'jpeg', 'webp'] and img.shape[-1] >= 3:
|
||||
img = cv.cvtColor(img, cv.COLOR_RGB2BGR)
|
||||
if img.shape[-1] >= 3:
|
||||
# ITU-R BT.709 luminance weights
|
||||
weights = np.array([0.2126, 0.7152, 0.0722])
|
||||
gray = np.dot(img[..., :3], weights)
|
||||
return gray[..., np.newaxis]
|
||||
|
||||
return img
|
||||
return img[..., 0:1] # Fallback for 2-channel images
|
||||
|
||||
def _prepare_image_data(self, img_tensor: torch.Tensor, file_type: str, bit_depth: int,
|
||||
save_as_grayscale: bool) -> np.ndarray:
|
||||
"""Prepare image data for saving by handling conversion and preprocessing."""
|
||||
# Convert from torch tensor if needed
|
||||
if isinstance(img_tensor, torch.Tensor):
|
||||
# Remove batch dimension if present
|
||||
if img_tensor.ndim == 4 and img_tensor.shape[0] == 1:
|
||||
img_np = img_tensor.squeeze(0).cpu().numpy()
|
||||
else:
|
||||
img_np = img_tensor.cpu().numpy()
|
||||
def prepare_image(self, img_tensor: torch.Tensor, save_as_grayscale: bool) -> np.ndarray:
|
||||
"""Convert tensor to numpy and prepare channels"""
|
||||
# Convert from tensor
|
||||
if img_tensor.ndim == 4 and img_tensor.shape[0] == 1:
|
||||
img_np = img_tensor.squeeze(0).cpu().numpy()
|
||||
else:
|
||||
img_np = img_tensor
|
||||
img_np = img_tensor.cpu().numpy()
|
||||
|
||||
sys.stderr.write(f"Initial data - Shape: {img_np.shape}, Min: {img_np.min()}, Max: {img_np.max()}\n")
|
||||
# Ensure float32 [0,1] range
|
||||
img_np = np.clip(img_np, 0, 1).astype(np.float32)
|
||||
|
||||
# Prepare image (handles channel conversion if needed)
|
||||
img_np = self._prepare_image_for_saving(img_np, file_type, save_as_grayscale)
|
||||
# Handle channels
|
||||
if len(img_np.shape) == 2:
|
||||
img_np = img_np[..., np.newaxis]
|
||||
|
||||
# Convert bit depth for all formats
|
||||
img_np = self._convert_bit_depth(img_np, bit_depth)
|
||||
sys.stderr.write(f"After bit_depth conversion min: {img_np.min()}, max: {img_np.max()}\n")
|
||||
# Convert to grayscale if explicitly requested (don't auto-detect for solid colors)
|
||||
if save_as_grayscale:
|
||||
img_np = self.convert_to_grayscale(img_np)
|
||||
|
||||
return img_np
|
||||
|
||||
def _save_exr(self, img_np: np.ndarray, out_path: str, save_as_grayscale: bool, exr_compression: str, bit_depth: int) -> None:
|
||||
"""Save image data as EXR file using OpenImageIO."""
|
||||
try:
|
||||
# Determine channels and prepare data
|
||||
if img_np.ndim == 2 or img_np.shape[-1] == 1 or save_as_grayscale:
|
||||
channels = 1
|
||||
if img_np.ndim == 3 and img_np.shape[-1] > 1:
|
||||
exr_data = img_np[..., 0:1] # Extract first channel for grayscale
|
||||
else:
|
||||
exr_data = img_np[..., np.newaxis] if img_np.ndim == 2 else img_np
|
||||
else:
|
||||
channels = img_np.shape[-1]
|
||||
exr_data = img_np
|
||||
|
||||
sys.stderr.write(f"\nSaving EXR - Shape: {exr_data.shape}, Channels: {channels}\n")
|
||||
sys.stderr.write(f"Value range - Min: {exr_data.min()}, Max: {exr_data.max()}\n")
|
||||
|
||||
# Ensure data is contiguous and in the right format
|
||||
if bit_depth == 16:
|
||||
exr_data = np.ascontiguousarray(exr_data.astype(np.float16))
|
||||
pixel_type = oiio.HALF
|
||||
else: # 32-bit
|
||||
exr_data = np.ascontiguousarray(exr_data.astype(np.float32))
|
||||
pixel_type = oiio.FLOAT
|
||||
|
||||
# Create spec with detected channels
|
||||
spec = oiio.ImageSpec(
|
||||
exr_data.shape[1], # width
|
||||
exr_data.shape[0], # height
|
||||
channels, # channels
|
||||
pixel_type # Use appropriate pixel type based on bit depth
|
||||
)
|
||||
|
||||
# Set EXR-specific attributes
|
||||
spec.attribute("compression", exr_compression)
|
||||
spec.attribute("Orientation", 1)
|
||||
spec.attribute("DateTime", datetime.now().isoformat())
|
||||
spec.attribute("Software", "COCO Tools")
|
||||
|
||||
# Create and write image buffer
|
||||
buf = oiio.ImageBuf(spec)
|
||||
buf.set_pixels(oiio.ROI(), exr_data)
|
||||
|
||||
if not buf.write(out_path):
|
||||
raise RuntimeError(f"Failed to write EXR: {oiio.geterror()}")
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"OpenImageIO EXR save failed: {str(e)}")
|
||||
|
||||
def _save_png(self, img_np: np.ndarray, out_path: str, bit_depth: int) -> None:
|
||||
"""Save image data as PNG file using OpenImageIO."""
|
||||
try:
|
||||
# Determine data type based on bit depth
|
||||
if bit_depth == 8:
|
||||
dtype = "uint8"
|
||||
pixel_type = np.uint8
|
||||
elif bit_depth == 16:
|
||||
dtype = "uint16"
|
||||
pixel_type = np.uint16
|
||||
else: # 32-bit
|
||||
dtype = "float"
|
||||
pixel_type = np.float32
|
||||
|
||||
# Convert and validate array
|
||||
png_data = np.ascontiguousarray(img_np.astype(pixel_type))
|
||||
|
||||
# Create image spec
|
||||
channels = 1 if png_data.ndim == 2 else png_data.shape[-1]
|
||||
spec = oiio.ImageSpec(
|
||||
png_data.shape[1], # width
|
||||
png_data.shape[0], # height
|
||||
channels, # channels
|
||||
dtype
|
||||
)
|
||||
|
||||
# Set PNG-specific attributes
|
||||
spec.attribute("compression", "zip")
|
||||
spec.attribute("png:compressionLevel", 9)
|
||||
|
||||
if bit_depth == 32:
|
||||
spec.attribute("oiio:ColorSpace", "Linear")
|
||||
|
||||
# Create and write image buffer
|
||||
buf = oiio.ImageBuf(spec)
|
||||
buf.set_pixels(oiio.ROI(), png_data)
|
||||
|
||||
if not buf.write(out_path):
|
||||
raise RuntimeError(f"Failed to write PNG: {oiio.geterror()}")
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"OpenImageIO PNG save error: {str(e)}")
|
||||
|
||||
def _save_jpeg_or_webp(self, img_np: np.ndarray, out_path: str, quality: int) -> None:
|
||||
"""Save image data as JPEG or WebP file using OpenCV."""
|
||||
cv.imwrite(out_path, img_np, [cv.IMWRITE_JPEG_QUALITY, quality])
|
||||
|
||||
def _save_tiff(self, img_np: np.ndarray, out_path: str, bit_depth: int) -> None:
|
||||
"""Save image data as TIFF file using tifffile."""
|
||||
# Convert BGR back to RGB for TIFF saving
|
||||
if img_np.shape[-1] >= 3:
|
||||
img_np = cv.cvtColor(img_np, cv.COLOR_BGR2RGB)
|
||||
|
||||
# Handle TIFF saving with appropriate bit depth
|
||||
def save_exr(self, img: np.ndarray, path: str, bit_depth: int, compression: str) -> None:
|
||||
"""Save EXR with proper float handling"""
|
||||
# Convert to appropriate type based on bit depth
|
||||
if bit_depth == 8:
|
||||
tifffile.imwrite(out_path, img_np.astype(np.uint8), photometric='rgb')
|
||||
# 8-bit EXR as uint8
|
||||
data = (img * 255).astype(np.uint8)
|
||||
pixel_type = oiio.UINT8
|
||||
elif bit_depth == 16:
|
||||
tifffile.imwrite(out_path, img_np.astype(np.uint16), photometric='rgb')
|
||||
else: # 32-bit
|
||||
tifffile.imwrite(
|
||||
out_path,
|
||||
img_np.astype(np.float32),
|
||||
photometric='rgb',
|
||||
dtype=np.float32
|
||||
)
|
||||
data = img.astype(np.float16)
|
||||
pixel_type = oiio.HALF
|
||||
else: # 32
|
||||
data = img.astype(np.float32)
|
||||
pixel_type = oiio.FLOAT
|
||||
|
||||
data = np.ascontiguousarray(data)
|
||||
channels = 1 if data.ndim == 2 else data.shape[-1]
|
||||
|
||||
spec = oiio.ImageSpec(data.shape[1], data.shape[0], channels, pixel_type)
|
||||
spec.attribute("compression", compression)
|
||||
spec.attribute("Software", "COCO Tools")
|
||||
|
||||
buf = oiio.ImageBuf(spec)
|
||||
buf.set_pixels(oiio.ROI(), data)
|
||||
|
||||
if not buf.write(path):
|
||||
raise RuntimeError(f"Failed to write EXR: {oiio.geterror()}")
|
||||
|
||||
def save_images(self, images, file_path, file_type, bit_depth,
|
||||
quality=95, save_as_grayscale=False,
|
||||
version=1, start_frame=1001, frame_pad=4,
|
||||
prompt=None, extra_pnginfo=None, exr_compression="zips"):
|
||||
"""Save a batch of images in various formats with support for versioning and frame numbering."""
|
||||
def save_png(self, img: np.ndarray, path: str, bit_depth: int) -> None:
|
||||
"""Save PNG with proper bit depth handling"""
|
||||
# Convert to appropriate type based on bit depth
|
||||
if bit_depth == 8:
|
||||
data = (img * 255).astype(np.uint8)
|
||||
pixel_type = oiio.UINT8
|
||||
elif bit_depth == 16:
|
||||
data = (img * 65535).astype(np.uint16)
|
||||
pixel_type = oiio.UINT16
|
||||
elif bit_depth == 32:
|
||||
# 32-bit PNG as float32 (preserves full range)
|
||||
data = img.astype(np.float32)
|
||||
pixel_type = oiio.FLOAT
|
||||
else:
|
||||
# Fallback to 8-bit for unsupported depths
|
||||
data = (img * 255).astype(np.uint8)
|
||||
pixel_type = oiio.UINT8
|
||||
|
||||
data = np.ascontiguousarray(data)
|
||||
channels = 1 if data.ndim == 2 else data.shape[-1]
|
||||
|
||||
spec = oiio.ImageSpec(data.shape[1], data.shape[0], channels, pixel_type)
|
||||
spec.attribute("compression", "zip")
|
||||
spec.attribute("png:compressionLevel", 9)
|
||||
|
||||
# For 32-bit float PNG, set linear color space
|
||||
if bit_depth == 32:
|
||||
spec.attribute("oiio:ColorSpace", "Linear")
|
||||
|
||||
buf = oiio.ImageBuf(spec)
|
||||
buf.set_pixels(oiio.ROI(), data)
|
||||
|
||||
if not buf.write(path):
|
||||
raise RuntimeError(f"Failed to write PNG: {oiio.geterror()}")
|
||||
|
||||
def save_opencv_format(self, img: np.ndarray, path: str, quality: int = 95) -> None:
|
||||
"""Save JPEG/WebP using OpenCV"""
|
||||
# Convert to 8-bit BGR
|
||||
data = (img * 255).astype(np.uint8)
|
||||
|
||||
# Convert RGB to BGR only for 3+ channel images
|
||||
if data.shape[-1] >= 3:
|
||||
data = cv.cvtColor(data, cv.COLOR_RGB2BGR)
|
||||
|
||||
# Save with quality setting
|
||||
if path.endswith(('.jpg', '.jpeg')):
|
||||
cv.imwrite(path, data, [cv.IMWRITE_JPEG_QUALITY, quality])
|
||||
else: # webp
|
||||
cv.imwrite(path, data, [cv.IMWRITE_WEBP_QUALITY, quality])
|
||||
|
||||
def save_tiff(self, img: np.ndarray, path: str, bit_depth: int) -> None:
|
||||
"""Save TIFF with proper bit depth"""
|
||||
if bit_depth == 8:
|
||||
data = (img * 255).astype(np.uint8)
|
||||
elif bit_depth == 16:
|
||||
data = (img * 65535).astype(np.uint16)
|
||||
else: # 32
|
||||
data = img.astype(np.float32)
|
||||
|
||||
# Determine photometric interpretation
|
||||
photometric = 'minisblack' if data.shape[-1] == 1 else 'rgb'
|
||||
|
||||
tifffile.imwrite(path, data, photometric=photometric)
|
||||
|
||||
def get_unique_filepath(self, base_path: str) -> str:
|
||||
"""Get unique filepath with incremental counter"""
|
||||
if not os.path.exists(base_path):
|
||||
return base_path
|
||||
|
||||
dir_name = os.path.dirname(base_path)
|
||||
base_name = os.path.basename(base_path)
|
||||
name, ext = os.path.splitext(base_name)
|
||||
|
||||
counter = 1
|
||||
while True:
|
||||
new_path = os.path.join(dir_name, f"{name}_{counter}{ext}")
|
||||
if not os.path.exists(new_path):
|
||||
return new_path
|
||||
counter += 1
|
||||
|
||||
def save_images(self, images, file_path, filename, file_type, bit_depth,
|
||||
quality=95, save_as_grayscale=False, use_versioning=True,
|
||||
version=1, prompt=None, extra_pnginfo=None, exr_compression="zips"):
|
||||
"""Main save function with optimized pipeline"""
|
||||
try:
|
||||
# Validate inputs
|
||||
bit_depth = int(bit_depth)
|
||||
file_type, bit_depth = self._validate_format_bitdepth(file_type.lower(), bit_depth)
|
||||
file_ext = f".{file_type}"
|
||||
|
||||
# Handle absolute vs relative paths
|
||||
if not os.path.isabs(file_path):
|
||||
full_output_folder, filename, _, subfolder, _ = folder_paths.get_save_image_path(
|
||||
file_path, self.output_dir, images.shape[2], images.shape[3]
|
||||
)
|
||||
base_path = os.path.join(full_output_folder, filename)
|
||||
else:
|
||||
base_path = file_path
|
||||
os.makedirs(os.path.dirname(base_path), exist_ok=True)
|
||||
|
||||
# Handle versioning
|
||||
version_str = f"_v{version:03}" if version >= 0 else ""
|
||||
file_type = file_type.lower()
|
||||
bit_depth = self.validate_bit_depth(file_type, bit_depth)
|
||||
|
||||
# Process each image in the batch
|
||||
# Build base path
|
||||
if file_path:
|
||||
full_path = os.path.join(self.output_dir, file_path) if not os.path.isabs(file_path) else file_path
|
||||
os.makedirs(full_path, exist_ok=True)
|
||||
base_path = os.path.join(full_path, filename)
|
||||
else:
|
||||
base_path = os.path.join(self.output_dir, filename)
|
||||
|
||||
# Add version string
|
||||
version_str = f"_v{version:03d}" if use_versioning and version >= 0 else ""
|
||||
|
||||
# Process each image
|
||||
for i, img_tensor in enumerate(images):
|
||||
# Prepare image data
|
||||
img_np = self._prepare_image_data(img_tensor, file_type, bit_depth, save_as_grayscale)
|
||||
# Prepare image (all formats start from float32 [0,1])
|
||||
img_np = self.prepare_image(img_tensor, save_as_grayscale)
|
||||
|
||||
# Generate output filename
|
||||
frame_num = f".{str(start_frame + i).zfill(frame_pad)}" if file_type == "exr" else f"_{i:05d}"
|
||||
out_path = f"{base_path}{version_str}{frame_num}{file_ext}"
|
||||
# Build output path
|
||||
frame_str = f"_{i}" if len(images) > 1 else ""
|
||||
out_path = f"{base_path}{version_str}{frame_str}.{file_type}"
|
||||
out_path = self.get_unique_filepath(out_path)
|
||||
|
||||
# Use improved filename increment if file exists
|
||||
if os.path.exists(out_path):
|
||||
out_path = self.increment_filename(out_path)
|
||||
|
||||
# Save the image based on format
|
||||
# Save based on format
|
||||
if file_type == "exr":
|
||||
self._save_exr(img_np, out_path, save_as_grayscale, exr_compression, bit_depth)
|
||||
self.save_exr(img_np, out_path, bit_depth, exr_compression)
|
||||
elif file_type == "png":
|
||||
self._save_png(img_np, out_path, bit_depth)
|
||||
self.save_png(img_np, out_path, bit_depth)
|
||||
elif file_type in ["jpg", "jpeg", "webp"]:
|
||||
self._save_jpeg_or_webp(img_np, out_path, quality)
|
||||
self.save_opencv_format(img_np, out_path, quality)
|
||||
elif file_type == "tiff":
|
||||
self._save_tiff(img_np, out_path, bit_depth)
|
||||
|
||||
# Return success
|
||||
self.save_tiff(img_np, out_path, bit_depth)
|
||||
|
||||
return {"ui": {"images": []}}
|
||||
|
||||
|
||||
except Exception as e:
|
||||
sys.stderr.write(f"Error saving images: {e}\n")
|
||||
return {"ui": {"error": str(e)}}
|
||||
|
||||
# Register the node
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"saver": saver,
|
||||
}
|
||||
|
||||
# Optional: Register display names
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"saver": "Image Saver"
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "cocotools_io"
|
||||
description = "Advanced image input and output: EXR, 32 bit support and more"
|
||||
version = "0.2.2"
|
||||
version = "0.3.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["torch>=1.10.0", "numpy>=1.21.0", "colour-science>=0.4.2", "Pillow>=9.0.0", "opencv-python>=4.5.0", "tifffile", "OpenImageIO>=2.4.13.0", "rich>=10.0.0"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user