206 lines
9.4 KiB
Python
206 lines
9.4 KiB
Python
"""
|
|
Batch processing utilities for COCO Tools nodes
|
|
Provides shared functionality for handling tensor batch operations
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
import logging
|
|
from typing import Tuple, Union, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from .debug_utils import debug_log, format_tensor_info
|
|
except ImportError:
|
|
# Use fallback functions
|
|
debug_log = lambda logger, level, simple_msg, verbose_msg=None, **kwargs: getattr(logger, level.lower())(simple_msg)
|
|
format_tensor_info = lambda tensor_shape, tensor_dtype, name="": f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}"
|
|
|
|
|
|
class BatchProcessor:
|
|
"""Utility class for batch tensor operations"""
|
|
|
|
@staticmethod
|
|
def validate_batch_tensor(tensor: torch.Tensor, expected_dims: int = 4,
|
|
tensor_name: str = "input") -> Tuple[int, int, int, int]:
|
|
"""
|
|
Validate and extract batch tensor dimensions
|
|
|
|
Args:
|
|
tensor: Input tensor to validate
|
|
expected_dims: Expected number of dimensions (default 4 for [B,H,W,C])
|
|
tensor_name: Name for error messages
|
|
|
|
Returns:
|
|
Tuple of (batch_size, height, width, channels)
|
|
|
|
Raises:
|
|
ValueError: If tensor doesn't match expected format
|
|
"""
|
|
if len(tensor.shape) != expected_dims:
|
|
raise ValueError(f"Expected {expected_dims}D {tensor_name} tensor, got {len(tensor.shape)}D with shape {tensor.shape}")
|
|
|
|
if expected_dims == 4:
|
|
batch_size, height, width, channels = tensor.shape
|
|
debug_log(logger, "debug", f"Validated {tensor_name}: B={batch_size}, H={height}, W={width}, C={channels}",
|
|
f"Validated {tensor_name} tensor: batch_size={batch_size}, height={height}, width={width}, channels={channels}")
|
|
return batch_size, height, width, channels
|
|
else:
|
|
raise NotImplementedError(f"Validation for {expected_dims}D tensors not implemented")
|
|
|
|
@staticmethod
|
|
def reshape_for_processing(tensor: torch.Tensor, preserve_alpha: bool = True) -> Tuple[np.ndarray, Optional[np.ndarray], Tuple]:
|
|
"""
|
|
Reshape batch tensor for element-wise processing (like colorspace conversion)
|
|
|
|
Args:
|
|
tensor: Input tensor [B, H, W, C]
|
|
preserve_alpha: Whether to separate alpha channel (for RGBA inputs)
|
|
|
|
Returns:
|
|
Tuple of (reshaped_rgb, alpha_channel_or_none, original_shape)
|
|
"""
|
|
# Convert to numpy and store original shape
|
|
img_np = tensor.cpu().numpy()
|
|
original_shape = img_np.shape
|
|
batch_size, height, width, channels = original_shape
|
|
|
|
debug_log(logger, "debug", f"Reshaping for processing: {format_tensor_info(original_shape, img_np.dtype)}",
|
|
f"Reshaping tensor from {original_shape} for batch processing")
|
|
|
|
# Reshape to [B*H*W, C] for element-wise processing
|
|
reshaped = img_np.reshape(-1, channels)
|
|
|
|
alpha_channel = None
|
|
if preserve_alpha and channels == 4:
|
|
# Separate alpha channel
|
|
alpha_channel = reshaped[..., 3:4]
|
|
rgb_data = reshaped[..., :3]
|
|
debug_log(logger, "debug", "Separated alpha channel",
|
|
f"Separated alpha channel: RGB shape={rgb_data.shape}, Alpha shape={alpha_channel.shape}")
|
|
return rgb_data, alpha_channel, original_shape
|
|
elif channels != 3 and not preserve_alpha:
|
|
# Handle non-RGB channels
|
|
if channels == 1:
|
|
# Replicate single channel to RGB
|
|
rgb_data = np.repeat(reshaped, 3, axis=-1)
|
|
debug_log(logger, "debug", "Replicated single channel to RGB",
|
|
f"Replicated single channel to RGB: {rgb_data.shape}")
|
|
elif channels > 3:
|
|
# Take first 3 channels
|
|
rgb_data = reshaped[..., :3]
|
|
debug_log(logger, "debug", f"Truncated to RGB from {channels} channels",
|
|
f"Truncated from {channels} channels to RGB: {rgb_data.shape}")
|
|
else:
|
|
# Pad to 3 channels
|
|
padding = np.zeros((reshaped.shape[0], 3 - channels))
|
|
rgb_data = np.concatenate([reshaped, padding], axis=-1)
|
|
debug_log(logger, "debug", f"Padded {channels} channels to RGB",
|
|
f"Padded from {channels} channels to RGB: {rgb_data.shape}")
|
|
return rgb_data, None, original_shape
|
|
else:
|
|
return reshaped, None, original_shape
|
|
|
|
@staticmethod
|
|
def reshape_from_processing(processed_rgb: np.ndarray, alpha_channel: Optional[np.ndarray],
|
|
original_shape: Tuple, target_device: torch.device) -> torch.Tensor:
|
|
"""
|
|
Reshape processed data back to original batch format
|
|
|
|
Args:
|
|
processed_rgb: Processed RGB data [B*H*W, 3]
|
|
alpha_channel: Optional alpha channel [B*H*W, 1]
|
|
original_shape: Original tensor shape (B, H, W, C)
|
|
target_device: Device to place result tensor on
|
|
|
|
Returns:
|
|
Reshaped tensor [B, H, W, C]
|
|
"""
|
|
batch_size, height, width, original_channels = original_shape
|
|
|
|
# Reshape RGB back to batch dimensions
|
|
rgb_reshaped = processed_rgb.reshape(batch_size, height, width, 3)
|
|
|
|
# Reattach alpha if present
|
|
if alpha_channel is not None:
|
|
alpha_reshaped = alpha_channel.reshape(batch_size, height, width, 1)
|
|
final_array = np.concatenate([rgb_reshaped, alpha_reshaped], axis=-1)
|
|
debug_log(logger, "debug", "Reattached alpha channel",
|
|
f"Reattached alpha channel: final shape={final_array.shape}")
|
|
else:
|
|
final_array = rgb_reshaped
|
|
|
|
# Convert back to torch tensor
|
|
result_tensor = torch.from_numpy(final_array).to(target_device)
|
|
|
|
debug_log(logger, "debug", f"Reshaped back to batch: {format_tensor_info(result_tensor.shape, result_tensor.dtype)}",
|
|
f"Reshaped processed data back to batch format: {result_tensor.shape}")
|
|
|
|
return result_tensor
|
|
|
|
@staticmethod
|
|
def normalize_batch_range(tensor: torch.Tensor, target_min: float = 0.0, target_max: float = 1.0,
|
|
source_min: Optional[float] = None, source_max: Optional[float] = None) -> torch.Tensor:
|
|
"""
|
|
Normalize batch tensor values to target range
|
|
|
|
Args:
|
|
tensor: Input tensor to normalize
|
|
target_min: Target minimum value
|
|
target_max: Target maximum value
|
|
source_min: Source minimum (if None, use tensor.min())
|
|
source_max: Source maximum (if None, use tensor.max())
|
|
|
|
Returns:
|
|
Normalized tensor
|
|
"""
|
|
if source_min is None:
|
|
source_min = tensor.min().item()
|
|
if source_max is None:
|
|
source_max = tensor.max().item()
|
|
|
|
debug_log(logger, "debug", f"Normalizing range: [{source_min:.6f}, {source_max:.6f}] -> [{target_min:.6f}, {target_max:.6f}]",
|
|
f"Batch normalization: source range [{source_min:.6f}, {source_max:.6f}] to target range [{target_min:.6f}, {target_max:.6f}]")
|
|
|
|
# Avoid division by zero
|
|
if source_max == source_min:
|
|
debug_log(logger, "warning", "Source range is zero, returning target_min",
|
|
f"Source min/max are equal ({source_min}), returning constant value {target_min}")
|
|
return torch.full_like(tensor, target_min)
|
|
|
|
# Normalize to [0,1] then scale to target range
|
|
normalized = (tensor - source_min) / (source_max - source_min)
|
|
scaled = normalized * (target_max - target_min) + target_min
|
|
|
|
return scaled
|
|
|
|
@staticmethod
|
|
def log_batch_info(tensor: torch.Tensor, operation: str, tensor_name: str = "tensor"):
|
|
"""
|
|
Log information about a batch tensor
|
|
|
|
Args:
|
|
tensor: Tensor to log info about
|
|
operation: Operation being performed
|
|
tensor_name: Name of the tensor for logging
|
|
"""
|
|
if len(tensor.shape) == 4:
|
|
batch_size = tensor.shape[0]
|
|
debug_log(logger, "info", f"{operation}: {batch_size} images, {format_tensor_info(tensor.shape, tensor.dtype, tensor_name)}",
|
|
f"{operation} - Batch processing {batch_size} images: {tensor_name} {format_tensor_info(tensor.shape, tensor.dtype)} " +
|
|
f"range=[{tensor.min().item():.6f}, {tensor.max().item():.6f}]")
|
|
else:
|
|
debug_log(logger, "info", f"{operation}: {format_tensor_info(tensor.shape, tensor.dtype, tensor_name)}",
|
|
f"{operation} - Processing {tensor_name}: {format_tensor_info(tensor.shape, tensor.dtype)} " +
|
|
f"range=[{tensor.min().item():.6f}, {tensor.max().item():.6f}]")
|
|
|
|
|
|
# Convenience functions for common operations
|
|
def validate_4d_batch(tensor: torch.Tensor, name: str = "input") -> Tuple[int, int, int, int]:
|
|
"""Convenience function for 4D batch validation"""
|
|
return BatchProcessor.validate_batch_tensor(tensor, 4, name)
|
|
|
|
def log_batch_processing(tensor: torch.Tensor, operation: str, name: str = "tensor"):
|
|
"""Convenience function for batch logging"""
|
|
BatchProcessor.log_batch_info(tensor, operation, name) |