CoCoandClaude Opus 4.6
62b62d2a96
- Add ruff configuration to pyproject.toml (py310, 120 char lines) - Fix trailing whitespace (W291/W293), modernize type annotations (UP006/UP045) - Remove unused imports (F401), simplify dict iteration (SIM118) - Use f-string conversion flags (RUF010), modernize exception aliases (UP024) - Replace try/except pass with contextlib.suppress (SIM105) - Prefix unused variables with underscore (F841/RUF059) - Convert lambdas to def statements (E731) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
375 lines
16 KiB
Python
375 lines
16 KiB
Python
import contextlib
|
|
import torch
|
|
import numpy as np
|
|
import colour
|
|
import logging
|
|
|
|
# Import batch processing utilities
|
|
try:
|
|
from ..utils.batch_utils import BatchProcessor, log_batch_processing
|
|
except ImportError:
|
|
# Fallback if utils not available
|
|
class BatchProcessor:
|
|
@staticmethod
|
|
def validate_batch_tensor(tensor, expected_dims=4, tensor_name="input"):
|
|
return tensor.shape
|
|
@staticmethod
|
|
def reshape_for_processing(tensor, preserve_alpha=True):
|
|
img_np = tensor.cpu().numpy()
|
|
return img_np.reshape(-1, img_np.shape[-1]), None, img_np.shape
|
|
@staticmethod
|
|
def reshape_from_processing(processed_rgb, alpha_channel, original_shape, target_device):
|
|
return torch.from_numpy(processed_rgb.reshape(original_shape)).to(target_device)
|
|
def log_batch_processing(tensor, operation, name="tensor"):
|
|
pass
|
|
|
|
# Import centralized logging setup
|
|
try:
|
|
from ..utils.debug_utils import setup_logging
|
|
setup_logging()
|
|
except ImportError:
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ColorspaceNode:
|
|
"""Simplified colorspace converter using colour-science library."""
|
|
|
|
# Map user-friendly names to colour-science names
|
|
colorspace_mapping = {
|
|
# ACES colorspaces
|
|
"ACES2065-1": "ACES2065-1", # Linear scene-referred ACES
|
|
"ACEScg": "ACEScg", # Linear scene-referred ACEScg
|
|
"ACEScct": "ACEScct", # Log-encoded ACEScg with toe
|
|
"ACEScc": "ACEScc", # Log-encoded ACES
|
|
|
|
# sRGB and Rec.709
|
|
"sRGB": "sRGB", # Standard display-referred sRGB (non-linear)
|
|
"sRGB Linear": "sRGB", # Linear version of sRGB
|
|
"Rec.709": "ITU-R BT.709", # Standard Rec.709 (scene-referred)
|
|
"Rec.709 Linear": "ITU-R BT.709", # Linear version of Rec.709
|
|
|
|
# Display P3
|
|
"Display P3": "Display P3", # Apple's Display P3 (non-linear)
|
|
"Display P3 Linear": "Display P3", # Linear version of Display P3
|
|
|
|
# Rec.2020
|
|
"Rec.2020": "ITU-R BT.2020", # Standard Rec.2020 (non-linear)
|
|
"Rec.2020 Linear": "ITU-R BT.2020", # Linear version of Rec.2020
|
|
|
|
# Adobe RGB
|
|
"Adobe RGB": "Adobe RGB (1998)", # Standard Adobe RGB (non-linear)
|
|
"Adobe RGB Linear": "Adobe RGB (1998)", # Linear version of Adobe RGB
|
|
|
|
# Raw/passthrough
|
|
"Raw": "Raw", # No colorspace conversion
|
|
}
|
|
|
|
# Track which colorspaces need encoding/decoding
|
|
encoded_spaces = {
|
|
"sRGB", # Standard sRGB is encoded (non-linear)
|
|
"ACEScct", # ACEScct is an encoded version of ACEScg
|
|
"ACEScc", # ACEScc is an encoded version of ACES2065-1
|
|
}
|
|
|
|
# Available colorspaces for the UI
|
|
available_colorspaces = list(colorspace_mapping.keys())
|
|
|
|
def __init__(self):
|
|
self.type = "ColorspaceNode"
|
|
logger.debug(f"Initialized with {len(self.available_colorspaces)} colorspaces")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define input types."""
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"from_colorspace": (cls.available_colorspaces,),
|
|
"to_colorspace": (cls.available_colorspaces,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "convert_colorspace"
|
|
CATEGORY = "COCO Tools/Processing"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, images, from_colorspace, to_colorspace, **kwargs):
|
|
"""
|
|
Smart caching based on input tensor properties and colorspace parameters.
|
|
Only re-execute when inputs or parameters actually change.
|
|
"""
|
|
try:
|
|
# Hash based on tensor properties
|
|
tensor_hash = f"{images.shape}_{images.dtype}_{images.device}"
|
|
|
|
# Include colorspace parameters
|
|
param_hash = f"{from_colorspace}_{to_colorspace}"
|
|
|
|
# Sample a few pixel values for content-aware caching
|
|
# Take samples from corners and center to detect actual data changes
|
|
sample_indices = [
|
|
(0, 0, 0, 0), # Top-left of first image
|
|
(0, -1, -1, 0), # Bottom-right of first image
|
|
(0, images.shape[1]//2, images.shape[2]//2, 0), # Center of first image
|
|
]
|
|
samples = []
|
|
for idx in sample_indices:
|
|
with contextlib.suppress(IndexError, RuntimeError):
|
|
samples.append(f"{images[idx].item():.6f}")
|
|
sample_hash = "_".join(samples)
|
|
|
|
return f"{tensor_hash}_{param_hash}_{sample_hash}"
|
|
except Exception:
|
|
# If anything goes wrong, always execute
|
|
return float("NaN")
|
|
|
|
def _is_encoded_colorspace(self, colorspace_name: str) -> bool:
|
|
"""Check if a colorspace name indicates encoded (non-linear) data."""
|
|
# Explicitly encoded spaces
|
|
if colorspace_name in self.encoded_spaces:
|
|
return True
|
|
|
|
# Linear spaces are not encoded
|
|
if "Linear" in colorspace_name:
|
|
return False
|
|
|
|
# Default encoding status for common colorspaces
|
|
default_encoded = {
|
|
"sRGB": True,
|
|
"Rec.709": True,
|
|
"Display P3": True,
|
|
"Rec.2020": True,
|
|
"Adobe RGB": True,
|
|
"ACES2065-1": False, # Linear by default
|
|
"ACEScg": False, # Linear by default
|
|
"Raw": False # Raw is always linear
|
|
}
|
|
|
|
# Check if the colorspace is in the default encoding status dictionary
|
|
if colorspace_name in default_encoded:
|
|
return default_encoded[colorspace_name]
|
|
|
|
# For unknown colorspaces, assume they are linear
|
|
return False
|
|
|
|
def _apply_gamma_encoding(self, rgb: np.ndarray, colorspace: str) -> np.ndarray:
|
|
"""Apply gamma encoding (linear → encoded) using inverse EOTF."""
|
|
# Handle standard colorspaces
|
|
if (colorspace == "sRGB" or "sRGB" in colorspace) and "Linear" not in colorspace:
|
|
# Apply inverse sRGB EOTF (linear → encoded)
|
|
return colour.models.eotf_inverse_sRGB(rgb)
|
|
elif (colorspace == "Rec.709" or "Rec.709" in colorspace) and "Linear" not in colorspace:
|
|
# BT.709 has its own OETF distinct from sRGB
|
|
return colour.models.oetf_BT709(rgb)
|
|
elif (colorspace == "Display P3" or "Display P3" in colorspace) and "Linear" not in colorspace:
|
|
# Display P3 uses the same transfer function as sRGB
|
|
return colour.models.eotf_inverse_sRGB(rgb)
|
|
elif (colorspace == "Rec.2020" or "Rec.2020" in colorspace) and "Linear" not in colorspace:
|
|
# Rec.2020 uses BT.1886 transfer function
|
|
return colour.models.eotf_inverse_BT1886(rgb)
|
|
elif (colorspace == "Adobe RGB" or "Adobe RGB" in colorspace) and "Linear" not in colorspace:
|
|
# Adobe RGB uses a gamma of 2.2
|
|
return np.power(np.maximum(rgb, 0), 1/2.2)
|
|
|
|
# Handle ACES colorspaces
|
|
elif colorspace == "ACEScc":
|
|
# Apply ACEScc encoding (log encoding for ACES)
|
|
return colour.models.log_encoding_ACEScc(rgb)
|
|
elif colorspace == "ACEScct":
|
|
# Apply ACEScct encoding (log encoding with toe for ACEScg)
|
|
return colour.models.log_encoding_ACEScct(rgb)
|
|
|
|
# Handle other gamma values
|
|
elif "Gamma 2.2" in colorspace:
|
|
return np.power(np.maximum(rgb, 0), 1/2.2)
|
|
elif "Gamma 2.4" in colorspace:
|
|
return np.power(np.maximum(rgb, 0), 1/2.4)
|
|
|
|
# Linear colorspaces don't need encoding
|
|
else:
|
|
return rgb
|
|
|
|
def _apply_gamma_decoding(self, rgb: np.ndarray, colorspace: str) -> np.ndarray:
|
|
"""Apply gamma decoding (encoded → linear) using EOTF."""
|
|
# Handle standard colorspaces
|
|
if (colorspace == "sRGB" or "sRGB" in colorspace) and "Linear" not in colorspace:
|
|
# Apply sRGB EOTF (encoded → linear)
|
|
return colour.models.eotf_sRGB(rgb)
|
|
elif (colorspace == "Rec.709" or "Rec.709" in colorspace) and "Linear" not in colorspace:
|
|
# BT.709 has its own OETF distinct from sRGB
|
|
return colour.models.oetf_inverse_BT709(rgb)
|
|
elif (colorspace == "Display P3" or "Display P3" in colorspace) and "Linear" not in colorspace:
|
|
# Display P3 uses the same transfer function as sRGB
|
|
return colour.models.eotf_sRGB(rgb)
|
|
elif (colorspace == "Rec.2020" or "Rec.2020" in colorspace) and "Linear" not in colorspace:
|
|
# Rec.2020 uses BT.1886 transfer function
|
|
return colour.models.eotf_BT1886(rgb)
|
|
elif (colorspace == "Adobe RGB" or "Adobe RGB" in colorspace) and "Linear" not in colorspace:
|
|
# Adobe RGB uses a gamma of 2.2
|
|
return np.power(np.maximum(rgb, 0), 2.2)
|
|
|
|
# Handle ACES colorspaces
|
|
elif colorspace == "ACEScc":
|
|
# Apply ACEScc decoding (inverse log encoding for ACES)
|
|
return colour.models.log_decoding_ACEScc(rgb)
|
|
elif colorspace == "ACEScct":
|
|
# Apply ACEScct decoding (inverse log encoding with toe for ACEScg)
|
|
return colour.models.log_decoding_ACEScct(rgb)
|
|
|
|
# Handle other gamma values
|
|
elif "Gamma 2.2" in colorspace:
|
|
return np.power(np.maximum(rgb, 0), 2.2)
|
|
elif "Gamma 2.4" in colorspace:
|
|
return np.power(np.maximum(rgb, 0), 2.4)
|
|
|
|
# Linear colorspaces don't need decoding
|
|
else:
|
|
return rgb
|
|
|
|
def convert_colorspace(self, images: torch.Tensor, from_colorspace: str, to_colorspace: str) -> tuple[torch.Tensor]:
|
|
"""
|
|
Convert images between colorspaces using colour-science library.
|
|
|
|
Args:
|
|
images: Input images as torch tensor [B, H, W, C]
|
|
from_colorspace: Source colorspace name
|
|
to_colorspace: Target colorspace name
|
|
|
|
Returns:
|
|
Tuple containing the converted images as torch tensor
|
|
"""
|
|
# Log batch processing info
|
|
log_batch_processing(images, f"Converting from '{from_colorspace}' to '{to_colorspace}'", "input")
|
|
|
|
# If source and target are the same, return original
|
|
if from_colorspace == to_colorspace:
|
|
logger.debug("Source and target colorspaces are the same")
|
|
return (images,)
|
|
|
|
# Convert to numpy
|
|
img_np = images.cpu().numpy()
|
|
logger.debug(f"Input range: min={img_np.min():.6f}, max={img_np.max():.6f}")
|
|
|
|
# Handle problematic values
|
|
if np.isnan(img_np).any() or np.isinf(img_np).any():
|
|
logger.warning("Input contains NaN/Inf values, cleaning...")
|
|
img_np = np.nan_to_num(img_np, nan=0.0, posinf=1.0, neginf=0.0)
|
|
images = torch.from_numpy(img_np).to(images.device)
|
|
|
|
try:
|
|
# Handle special cases first
|
|
if from_colorspace == "Raw" or to_colorspace == "Raw":
|
|
logger.debug("Raw colorspace detected, returning input unchanged")
|
|
return (images,)
|
|
|
|
# Get the colour-science colorspace names
|
|
from_cs = self.colorspace_mapping.get(from_colorspace)
|
|
to_cs = self.colorspace_mapping.get(to_colorspace)
|
|
|
|
if not from_cs or not to_cs:
|
|
logger.error(f"Unsupported colorspaces: {from_colorspace} -> {to_colorspace}")
|
|
return (images,)
|
|
|
|
# Handle encoding/decoding
|
|
working_img = img_np.copy()
|
|
|
|
# Step 1: Decode input if it's encoded
|
|
if self._is_encoded_colorspace(from_colorspace):
|
|
logger.debug(f"Decoding {from_colorspace}")
|
|
working_img = self._apply_gamma_decoding(working_img, from_colorspace)
|
|
|
|
# Step 2: Convert between colorspaces (linear to linear)
|
|
if from_cs != to_cs and from_cs != "Raw" and to_cs != "Raw":
|
|
logger.debug(f"Converting colorspace: {from_cs} -> {to_cs}")
|
|
|
|
# Handle special case where both map to same underlying space
|
|
if from_cs == to_cs:
|
|
logger.debug("Same underlying colorspace, skipping conversion")
|
|
else:
|
|
# Use batch processing utilities for reshaping
|
|
working_tensor = torch.from_numpy(working_img)
|
|
rgb_data, alpha_channel, original_shape = BatchProcessor.reshape_for_processing(
|
|
working_tensor, preserve_alpha=True
|
|
)
|
|
|
|
# Apply the colorspace conversion
|
|
try:
|
|
converted_rgb = colour.RGB_to_RGB(
|
|
rgb_data,
|
|
input_colourspace=from_cs,
|
|
output_colourspace=to_cs,
|
|
apply_cctf_decoding=False, # We handle encoding separately
|
|
apply_cctf_encoding=False
|
|
)
|
|
logger.debug("Colorspace conversion successful")
|
|
except Exception as e:
|
|
logger.error(f"Colour-science conversion failed: {e}")
|
|
# Try with chromatic adaptation
|
|
try:
|
|
converted_rgb = colour.RGB_to_RGB(
|
|
rgb_data,
|
|
input_colourspace=from_cs,
|
|
output_colourspace=to_cs,
|
|
apply_cctf_decoding=False,
|
|
apply_cctf_encoding=False,
|
|
chromatic_adaptation_transform='CAT02'
|
|
)
|
|
logger.debug("Colorspace conversion with CAT02 successful")
|
|
except Exception as e2:
|
|
logger.error(f"All conversion attempts failed: {e2}")
|
|
# Return original image
|
|
return (images,)
|
|
|
|
# Use batch processing utilities to reshape back
|
|
working_tensor = BatchProcessor.reshape_from_processing(
|
|
converted_rgb, alpha_channel, original_shape, images.device
|
|
)
|
|
working_img = working_tensor.cpu().numpy()
|
|
|
|
# Step 3: Encode output if needed
|
|
if self._is_encoded_colorspace(to_colorspace):
|
|
logger.debug(f"Encoding to {to_colorspace}")
|
|
working_img = self._apply_gamma_encoding(working_img, to_colorspace)
|
|
|
|
# Handle clipping based on colorspace
|
|
# For HDR colorspaces like ACES, we don't want to clip to 0-1
|
|
is_hdr_colorspace = any(hdr_space in to_colorspace for hdr_space in ["ACES", "Raw", "Linear"])
|
|
is_aces = any(aces in to_colorspace for aces in ["ACEScg", "ACES2065", "ACEScc", "ACEScct"])
|
|
|
|
if not is_hdr_colorspace:
|
|
# For display-referred spaces, clip to 0-1
|
|
working_img = np.clip(working_img, 0.0, 1.0)
|
|
elif not is_aces:
|
|
# For non-ACES HDR/scene-referred spaces, clamp negatives but allow values > 1.0
|
|
working_img = np.maximum(working_img, 0.0)
|
|
|
|
# Log if we have values > 1.0 (common in HDR)
|
|
if np.any(working_img > 1.0):
|
|
max_val = np.max(working_img)
|
|
logger.debug(f"HDR values detected: max={max_val:.6f}")
|
|
else:
|
|
# ACES colorspaces legitimately use negative values, preserve them
|
|
neg_count = np.sum(working_img < 0)
|
|
if neg_count > 0:
|
|
logger.debug(f"ACES colorspace preserving {neg_count} negative values")
|
|
|
|
# Log if we have values > 1.0 (common in HDR/ACES)
|
|
if np.any(working_img > 1.0):
|
|
max_val = np.max(working_img)
|
|
logger.debug(f"ACES HDR values detected: max={max_val:.6f}")
|
|
|
|
# Convert back to torch tensor
|
|
result_tensor = torch.from_numpy(working_img).to(images.device)
|
|
|
|
logger.debug(f"Output range: min={result_tensor.min().item():.6f}, max={result_tensor.max().item():.6f}")
|
|
|
|
return (result_tensor,)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Conversion failed: {e}")
|
|
import traceback
|
|
logger.error(traceback.format_exc())
|
|
return (images,)
|