nodes.py

import hashlib

import numpy as np
import torch
from PIL import Image
import cv2
import tempfile


def _is_changed_hash(image=None, **kwargs) -> str:
    """Deterministic hash over scalar node inputs for ComfyUI caching.

    Tensor inputs are summarised by shape/dtype; ComfyUI's graph already
    tracks upstream tensor changes via connections so hashing tensor
    content would be wasted work.
    """
    m = hashlib.sha256()
    for key in sorted(kwargs.keys()):
        m.update(f"{key}={kwargs[key]!r};".encode())
    if image is not None and hasattr(image, "shape"):
        m.update(f"image_shape={tuple(image.shape)};image_dtype={getattr(image, 'dtype', None)};".encode())
    return m.hexdigest()

# Handle ComfyUI imports gracefully for testing
try:
    import folder_paths
    import comfy.utils
    COMFYUI_AVAILABLE = True
except ImportError:
    # Mock modules for testing/CI
    import sys
    from types import ModuleType
    
    # Create mock folder_paths
    folder_paths = ModuleType('folder_paths')
    folder_paths.get_input_directory = lambda: tempfile.gettempdir()
    sys.modules['folder_paths'] = folder_paths
    
    # Create mock comfy.utils
    comfy = ModuleType('comfy')
    comfy.utils = ModuleType('comfy.utils')
    comfy.utils.common_upscale = lambda image, width, height, upscale_method, crop: image
    sys.modules['comfy'] = comfy
    sys.modules['comfy.utils'] = comfy.utils
    
    COMFYUI_AVAILABLE = False

class TransparencyBackgroundRemover:
    """
    ComfyUI node for automatic background removal with transparency generation.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "tolerance": ("INT", {
                    "default": 30,
                    "min": 0,
                    "max": 255,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Color similarity threshold for background detection (0-255)"
                }),
                "edge_sensitivity": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "slider",
                    "tooltip": "Edge detection sensitivity (0-1)"
                }),
                "foreground_bias": ("FLOAT", {
                    "default": 0.7,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "slider",
                    "tooltip": "Bias towards foreground preservation (0-1)"
                }),
                "color_clusters": ("INT", {
                    "default": 8,
                    "min": 2,
                    "max": 20,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Number of color clusters for background detection"
                }),
                "binary_threshold": ("INT", {
                    "default": 128,
                    "min": 0,
                    "max": 255,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Threshold for binary alpha mask (0-255)"
                }),
                "output_size": (["ORIGINAL", "64x64", "96x96", "128x128", "256x256", "512x512", "768x768", "1024x1024", "1280x1280", "1536x1536", "1792x1792", "2048x2048"], {
                    "default": "ORIGINAL",
                    "tooltip": "Target output size (power-of-8 dimensions for optimal scaling)"
                }),
                "scaling_method": (["NEAREST", "BILINEAR", "BICUBIC", "LANCZOS"], {
                    "default": "NEAREST",
                    "tooltip": "Interpolation method: NEAREST (pixel-perfect), BILINEAR (smooth), BICUBIC (high-quality), LANCZOS (best quality)"
                }),
            },
            "optional": {
                "edge_refinement": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Apply edge refinement post-processing"
                }),
                "dither_handling": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Enable dithered pattern detection and handling"
                }),
                "output_format": (["RGBA", "RGB_WITH_MASK"], {
                    "default": "RGBA",
                    "tooltip": "Output format: RGBA with alpha channel or RGB with separate mask"
                }),
                "auto_adjust": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Automatically adjust parameters based on image content analysis"
                }),
                "edge_detection_mode": (["AUTO", "PIXEL_ART", "PHOTOGRAPHIC"], {
                    "default": "AUTO",
                    "tooltip": "Edge detection optimization: AUTO (detect content type), PIXEL_ART (sharp edges), PHOTOGRAPHIC (smooth edges)"
                }),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "mask")
    FUNCTION = "remove_background"
    CATEGORY = "image/processing"
    OUTPUT_NODE = False
    DESCRIPTION = (
        "Transparency-aware background removal using multi-algorithm "
        "detection (edge, color clustering, corner sampling, dither). "
        "Returns an IMAGE (RGBA or RGB depending on output_format) and MASK."
    )

    @classmethod
    def IS_CHANGED(cls, image=None, **kwargs):
        return _is_changed_hash(image, **kwargs)

    def __init__(self):
        self.processor = None
    
    def parse_output_size(self, size_string):
        """
        Parse output size string to width, height tuple.
        
        Args:
            size_string: String like "512x512" or "ORIGINAL"
            
        Returns:
            Tuple (width, height) or None for ORIGINAL
        """
        if size_string == "ORIGINAL":
            return None
        
        try:
            width, height = size_string.split('x')
            return (int(width), int(height))
        except ValueError:
            raise ValueError(f"Invalid output size format: {size_string}")
    
    def calculate_scaling_factor(self, current_size, target_size):
        """
        Calculate optimal integer scaling factor for NEAREST interpolation.
        
        Args:
            current_size: Tuple (width, height) of current image
            target_size: Tuple (width, height) of target size
            
        Returns:
            Scaling factor that achieves target size
        """
        current_w, current_h = current_size
        target_w, target_h = target_size
        
        # Calculate scale factors for width and height
        scale_w = target_w / current_w
        scale_h = target_h / current_h
        
        # Use the same scale for both dimensions to maintain aspect ratio
        # Choose the smaller scale to ensure we don't exceed target dimensions
        scale_factor = min(scale_w, scale_h)
        
        return scale_factor
    
    def intelligent_scale(self, image_pil, target_size, scaling_method="NEAREST"):
        """
        Scale image to target dimensions using specified interpolation method.
        
        Args:
            image_pil: PIL Image object
            target_size: Tuple (width, height) for target dimensions
            scaling_method: Interpolation method ("NEAREST", "BILINEAR", "BICUBIC", "LANCZOS")
            
        Returns:
            Scaled PIL Image
        """
        if target_size is None:
            return image_pil
            
        current_size = (image_pil.width, image_pil.height)
        target_w, target_h = target_size
        
        # If already at target size, return as-is
        if current_size == target_size:
            return image_pil
        
        # Calculate scaling factor
        scale_factor = self.calculate_scaling_factor(current_size, target_size)
        
        # Apply scaling
        new_width = int(image_pil.width * scale_factor)
        new_height = int(image_pil.height * scale_factor)
        
        # If calculated size matches target exactly, use target dimensions
        if abs(new_width - target_w) <= 1 and abs(new_height - target_h) <= 1:
            new_width, new_height = target_w, target_h
        
        # Select resampling method based on scaling_method
        resampling_map = {
            "NEAREST": Image.Resampling.NEAREST,
            "BILINEAR": Image.Resampling.BILINEAR, 
            "BICUBIC": Image.Resampling.BICUBIC,
            "LANCZOS": Image.Resampling.LANCZOS
        }
        
        resampling_method = resampling_map.get(scaling_method, Image.Resampling.NEAREST)
        
        return image_pil.resize(
            (new_width, new_height),
            resampling_method
        )

    @torch.no_grad()
    def remove_background(self, image, tolerance=30, edge_sensitivity=0.8,
                         foreground_bias=0.7, color_clusters=8, binary_threshold=128,
                         edge_refinement=True, dither_handling=True, output_format="RGBA",
                         output_size="ORIGINAL", scaling_method="NEAREST", auto_adjust=False,
                         edge_detection_mode="AUTO"):
        """
        Main processing function for background removal with error handling.
        """
        try:
            # Validate input
            if image is None or image.shape[0] == 0:
                raise ValueError("No input image provided")

            # Check image dimensions
            if len(image.shape) != 4:
                raise ValueError(f"Expected 4D tensor, got {len(image.shape)}D")
            
            # Validate minimum input size (64x64 pixels)
            _, height, width, _ = image.shape
            if height < 64 or width < 64:
                raise ValueError(f"Input image must be at least 64x64 pixels, got {width}x{height}")

            # Process with error catching
            results, masks = self._process_images(
                image=image, 
                tolerance=tolerance,
                edge_sensitivity=edge_sensitivity,
                foreground_bias=foreground_bias,
                color_clusters=color_clusters,
                edge_refinement=edge_refinement,
                dither_handling=dither_handling,
                binary_threshold=binary_threshold,
                output_format=output_format,
                output_size=output_size,
                scaling_method=scaling_method,
                auto_adjust=auto_adjust,
                edge_detection_mode=edge_detection_mode
            )

            return (results, masks)

        except cv2.error as e:
            raise RuntimeError(f"OpenCV processing error: {str(e)}")
        except MemoryError:
            raise RuntimeError("Insufficient memory for processing. Try reducing batch size.")
        except Exception as e:
            raise RuntimeError(f"Background removal failed: {str(e)}")

    def _process_images(self, image, tolerance=30, edge_sensitivity=0.8,
                       foreground_bias=0.7, color_clusters=8, binary_threshold=128,
                       edge_refinement=True, dither_handling=True, output_format="RGBA",
                       output_size="ORIGINAL", scaling_method="NEAREST", auto_adjust=False,
                       edge_detection_mode="AUTO"):
        """
        Internal method for processing images without error handling wrapper.
        """
        # Convert from ComfyUI tensor format to numpy
        batch_size, height, width, channels = image.shape
        results = []
        masks = []

        for i in range(batch_size):
            # Convert single image to numpy array
            img_np = (image[i].cpu().numpy() * 255).astype(np.uint8)

            # Initialize processor with parameters
            try:
                from .background_remover import EnhancedPixelArtProcessor
            except ImportError:
                # Fallback for testing outside package structure
                from background_remover import EnhancedPixelArtProcessor
            
            # Determine dither handling based on edge detection mode
            effective_dither_handling = dither_handling
            if edge_detection_mode == "AUTO":
                # Use content detection for AUTO mode
                temp_processor = EnhancedPixelArtProcessor()
                is_pixel_art = temp_processor._detect_pixel_art_characteristics(img_np)
                effective_dither_handling = is_pixel_art
            elif edge_detection_mode == "PIXEL_ART":
                effective_dither_handling = True
            elif edge_detection_mode == "PHOTOGRAPHIC":
                effective_dither_handling = False
            
            processor = EnhancedPixelArtProcessor(
                tolerance=tolerance,
                edge_sensitivity=edge_sensitivity,
                color_clusters=color_clusters,
                foreground_bias=foreground_bias,
                edge_refinement=edge_refinement,
                dither_handling=effective_dither_handling,
                binary_threshold=binary_threshold
            )

            # Auto-adjust parameters if enabled
            if auto_adjust:
                adjustments = processor.auto_adjust_parameters(img_np)
                if adjustments:
                    # Apply adjustments
                    if 'tolerance' in adjustments:
                        processor.tolerance = adjustments['tolerance']
                    if 'edge_sensitivity' in adjustments:
                        processor.edge_sensitivity = adjustments['edge_sensitivity']
                    if 'foreground_bias' in adjustments:
                        processor.foreground_bias = adjustments['foreground_bias']

            # Process image
            rgba_result = processor.remove_background_advanced(img_np)
            
            # Apply scaling if requested
            if output_size != "ORIGINAL":
                # Parse target dimensions
                target_dimensions = self.parse_output_size(output_size)
                
                if target_dimensions is not None:
                    # Convert to PIL Image for scaling
                    rgba_pil = Image.fromarray(rgba_result, 'RGBA')
                    
                    # Apply intelligent scaling with specified method
                    rgba_scaled = self.intelligent_scale(rgba_pil, target_dimensions, scaling_method)
                    
                    # Convert back to numpy
                    rgba_result = np.array(rgba_scaled)

            # Extract alpha channel as mask (invert: 0=transparent, 255=opaque)
            alpha_channel = rgba_result[:, :, 3]
            masks.append(alpha_channel)

            # Handle output format
            if output_format == "RGBA":
                results.append(rgba_result)
            else:  # RGB_WITH_MASK
                rgb_result = rgba_result[:, :, :3]
                results.append(rgb_result)

        # Convert back to ComfyUI tensor format
        if output_format == "RGBA":
            # For RGBA, maintain 4 channels
            result_tensor = torch.from_numpy(np.array(results)).to(dtype=torch.float32) / 255.0
        else:
            # For RGB_WITH_MASK, use 3 channels
            result_tensor = torch.from_numpy(np.array(results)).to(dtype=torch.float32) / 255.0
        
        mask_tensor = torch.from_numpy(np.array(masks)).to(dtype=torch.float32) / 255.0

        return result_tensor, mask_tensor

class TransparencyBackgroundRemoverBatch:
    """
    Batch processing version with additional options and auto-adjustment.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "tolerance": ("INT", {
                    "default": 30,
                    "min": 0,
                    "max": 255,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Base color similarity threshold for background detection (0-255)"
                }),
                "edge_sensitivity": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "slider",
                    "tooltip": "Base edge detection sensitivity (0-1)"
                }),
                "auto_adjust": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Automatically adjust parameters based on image content"
                }),
                "foreground_bias": ("FLOAT", {
                    "default": 0.7,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "slider",
                    "tooltip": "Bias towards foreground preservation (0-1)"
                }),
                "color_clusters": ("INT", {
                    "default": 8,
                    "min": 2,
                    "max": 20,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Number of color clusters for background detection"
                }),
            },
            "optional": {
                "edge_refinement": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Apply edge refinement post-processing"
                }),
                "dither_handling": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Enable dithered pattern detection and handling"
                }),
                "binary_threshold": ("INT", {
                    "default": 128,
                    "min": 0,
                    "max": 255,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Threshold for binary alpha mask (0-255)"
                }),
                "progress_reporting": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Generate detailed processing report"
                }),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
    RETURN_NAMES = ("images", "masks", "report")
    FUNCTION = "batch_remove_background"
    CATEGORY = "image/processing"
    OUTPUT_NODE = False
    DESCRIPTION = (
        "Batch version of the transparency background remover with optional "
        "per-image parameter auto-adjustment. Produces a stacked IMAGE, MASK, "
        "and a STRING summary report."
    )

    @classmethod
    def IS_CHANGED(cls, images=None, **kwargs):
        return _is_changed_hash(images, **kwargs)

    @torch.no_grad()
    def batch_remove_background(self, images, tolerance=30, edge_sensitivity=0.8,
                               auto_adjust=True, foreground_bias=0.7, color_clusters=8,
                               edge_refinement=True, dither_handling=True, binary_threshold=128,
                               progress_reporting=True):
        """
        Batch process multiple images with optional auto-adjustment.
        
        Returns:
            tuple: A tuple containing:
                - result_tensor (torch.Tensor): RGBA images with transparent backgrounds (batch, height, width, 4)
                - mask_tensor (torch.Tensor): Binary alpha masks (batch, height, width, 1) 
                - summary_report (str): Detailed processing report with stats and any errors
        """
        try:
            if images is None or images.shape[0] == 0:
                raise ValueError("No input images provided")

            batch_size, height, width, channels = images.shape
            
            # Validate input dimensions
            if len(images.shape) != 4:
                raise ValueError(f"Expected 4D tensor (batch, height, width, channels), got {len(images.shape)}D")
            
            if channels not in [3, 4]:
                raise ValueError(f"Expected 3 or 4 channels (RGB or RGBA), got {channels}")
            
            # Validate minimum input size for all images
            if height < 64 or width < 64:
                raise ValueError(f"All input images must be at least 64x64 pixels, got {width}x{height}")

            results = []
            masks = []
            reports = []
            processing_stats = {
                'total_images': batch_size,
                'successful': 0,
                'failed': 0,
                'auto_adjustments': 0,
                'avg_processing_time': 0.0,
                'adjustments_made': []
            }

            import time
            total_processing_time = 0.0

            for i in range(batch_size):
                start_time = time.time()
                
                try:
                    # Convert single image to numpy array
                    img_np = (images[i].cpu().numpy() * 255).astype(np.uint8)

                    # Initialize processor with base parameters
                    try:
                        from .background_remover import EnhancedPixelArtProcessor
                    except ImportError:
                        # Fallback for testing outside package structure
                        from background_remover import EnhancedPixelArtProcessor
                    processor = EnhancedPixelArtProcessor(
                        tolerance=tolerance,
                        edge_sensitivity=edge_sensitivity,
                        color_clusters=color_clusters,
                        foreground_bias=foreground_bias,
                        edge_refinement=edge_refinement,
                        dither_handling=dither_handling,
                        binary_threshold=binary_threshold
                    )

                    # Auto-adjust parameters if enabled
                    adjustments = {}
                    if auto_adjust:
                        adjustments = processor.auto_adjust_parameters(img_np)
                        if adjustments:
                            processing_stats['auto_adjustments'] += 1
                            processing_stats['adjustments_made'].append({
                                'image_index': i,
                                'adjustments': adjustments
                            })
                            
                            # Apply adjustments
                            if 'tolerance' in adjustments:
                                processor.tolerance = adjustments['tolerance']
                            if 'edge_sensitivity' in adjustments:
                                processor.edge_sensitivity = adjustments['edge_sensitivity']
                            if 'foreground_bias' in adjustments:
                                processor.foreground_bias = adjustments['foreground_bias']

                    # Process image
                    rgba_result = processor.remove_background_advanced(img_np)

                    # Extract alpha channel as mask
                    alpha_channel = rgba_result[:, :, 3]
                    masks.append(alpha_channel)
                    results.append(rgba_result)

                    processing_stats['successful'] += 1
                    
                    processing_time = time.time() - start_time
                    total_processing_time += processing_time

                    if progress_reporting:
                        report = f"Image {i+1}: Processed successfully"
                        if adjustments:
                            report += f" (auto-adjusted: {list(adjustments.keys())})"
                        report += f" in {processing_time:.3f}s"
                        reports.append(report)

                except Exception as e:
                    processing_stats['failed'] += 1
                    error_msg = f"Image {i+1}: Failed - {str(e)}"
                    reports.append(error_msg)
                    
                    # Create empty result for failed image
                    empty_result = np.zeros_like(img_np if 'img_np' in locals() else (height, width, 4), dtype=np.uint8)
                    if len(empty_result.shape) == 3 and empty_result.shape[2] == 3:
                        empty_rgba = np.zeros((empty_result.shape[0], empty_result.shape[1], 4), dtype=np.uint8)
                        empty_rgba[:, :, :3] = empty_result
                        empty_result = empty_rgba
                    results.append(empty_result)
                    masks.append(np.zeros((height, width), dtype=np.uint8))

            # Calculate statistics
            processing_stats['avg_processing_time'] = total_processing_time / batch_size if batch_size > 0 else 0.0

            # Convert results to tensors
            result_tensor = torch.from_numpy(np.array(results)).to(dtype=torch.float32) / 255.0
            mask_tensor = torch.from_numpy(np.array(masks)).to(dtype=torch.float32) / 255.0

            # Generate summary report
            summary_report = self._generate_summary_report(processing_stats, reports if progress_reporting else [])

            return (result_tensor, mask_tensor, summary_report)

        except Exception as e:
            error_report = f"Batch processing failed: {str(e)}"
            # Return empty tensors in case of complete failure
            empty_images = torch.zeros_like(images)
            empty_masks = torch.zeros((images.shape[0], images.shape[1], images.shape[2]))
            return (empty_images, empty_masks, error_report)

    def _generate_summary_report(self, stats, detailed_reports):
        """Generate a summary report of batch processing results."""
        lines = [
            "=== Batch Processing Report ===",
            f"Total Images: {stats['total_images']}",
            f"Successful: {stats['successful']}",
            f"Failed: {stats['failed']}",
            f"Success Rate: {(stats['successful']/stats['total_images']*100):.1f}%" if stats['total_images'] > 0 else "N/A",
            f"Auto-adjustments Applied: {stats['auto_adjustments']}",
            f"Average Processing Time: {stats['avg_processing_time']:.3f}s per image",
            ""
        ]

        if stats['adjustments_made']:
            lines.append("Auto-adjustment Details:")
            for adj in stats['adjustments_made']:
                lines.append(f"  Image {adj['image_index']+1}: {adj['adjustments']}")
            lines.append("")

        if detailed_reports:
            lines.append("Detailed Processing Log:")
            lines.extend(detailed_reports)

        return "\n".join(lines)

# Node registration
NODE_CLASS_MAPPINGS = {
    "TransparencyBackgroundRemover": TransparencyBackgroundRemover,
    "TransparencyBackgroundRemoverBatch": TransparencyBackgroundRemoverBatch,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TransparencyBackgroundRemover": "Transparency Background Remover",
    "TransparencyBackgroundRemoverBatch": "Transparency Background Remover (Batch)",
}