diff --git a/__init__.py b/__init__.py index 9091b1a..bf6e72e 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/js/saver.js b/js/saver.js index d60d955..e7eb87a 100644 --- a/js/saver.js +++ b/js/saver.js @@ -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; }; } -}); \ No newline at end of file +}); diff --git a/modules/saver.py b/modules/saver.py index 1e2de2a..b06cf1f 100644 --- a/modules/saver.py +++ b/modules/saver.py @@ -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" -} \ No newline at end of file +} diff --git a/node.zip b/node.zip index af8f728..edcc2c8 100644 Binary files a/node.zip and b/node.zip differ diff --git a/pyproject.toml b/pyproject.toml index 6b67628..367b308 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"]