diff --git a/__init__.py b/__init__.py index 0a8e6dd..e8c54f2 100644 --- a/__init__.py +++ b/__init__.py @@ -491,6 +491,25 @@ class CropImageBatchToMaskBoundsNodeAE: Returns: tuple: (cropped_images, cropped_masks) """ + # Validate tensor shapes and batch sizes + if not isinstance(images, torch.Tensor): + raise ValueError(f"images must be a torch.Tensor, got {type(images).__name__}") + + if not isinstance(masks, torch.Tensor): + raise ValueError(f"masks must be a torch.Tensor, got {type(masks).__name__}") + + if len(images.shape) != 4: + raise ValueError(f"images must be 4D (B, H, W, C), got shape {images.shape}") + + if len(masks.shape) != 3: + raise ValueError(f"masks must be 3D (B, H, W), got shape {masks.shape}") + + if images.shape[0] != masks.shape[0]: + raise ValueError(f"images and masks must have same batch size, got images: {images.shape[0]}, masks: {masks.shape[0]}") + + if images.shape[1] != masks.shape[1] or images.shape[2] != masks.shape[2]: + raise ValueError(f"images and masks must have same spatial dimensions, got images: ({images.shape[1]}, {images.shape[2]}), masks: ({masks.shape[1]}, {masks.shape[2]})") + # Validate inputs if mode not in ["mask_for_each", "sum_of_masks"]: raise ValueError(f"mode must be 'mask_for_each' or 'sum_of_masks', got: {mode}") diff --git a/image_selector.py b/image_selector.py index 2271a8d..d208547 100644 --- a/image_selector.py +++ b/image_selector.py @@ -45,6 +45,10 @@ def create_thumbnail(img, size): new_height = size[1] new_width = int(size[1] * img_ratio) + # Ensure dimensions are at least 1 pixel (prevent resize errors with extreme aspect ratios) + new_width = max(1, new_width) + new_height = max(1, new_height) + return img.resize((new_width, new_height), Image.Resampling.LANCZOS) # Cache management @@ -276,7 +280,23 @@ def sort_images(images, paths, sort_mode, ascending=True): return sorted_images, sorted_paths def pil_to_tensor(img): - """Convert PIL image to torch tensor (1,H,W,3)""" + """Convert PIL image to torch tensor (1,H,W,3) + + Args: + img: PIL Image object + + Returns: + tuple: (image_tensor, mask_tensor) + + Raises: + ValueError: If img is None or not a PIL Image + """ + if img is None: + raise ValueError("Image cannot be None") + + if not hasattr(img, 'mode'): + raise ValueError(f"Expected PIL Image, got {type(img).__name__}") + if img.mode == 'RGBA': # Handle transparency img_rgb = img.convert('RGB') diff --git a/tests/test_crop_mask.py b/tests/test_crop_mask.py index 2ced3b3..d9db66b 100644 --- a/tests/test_crop_mask.py +++ b/tests/test_crop_mask.py @@ -274,3 +274,69 @@ def test_crop_node_batch_consistency(): for i in range(1, 3): assert result_images[i].shape == result_images[0].shape assert result_masks[i].shape == result_masks[0].shape + + +def test_crop_node_mismatched_batch_size(): + """Test error handling for mismatched batch sizes""" + from __init__ import CropImageBatchToMaskBoundsNodeAE + import torch + import pytest + + node = CropImageBatchToMaskBoundsNodeAE() + + # Mismatched batch sizes + images = torch.zeros(3, 64, 64, 3) + masks = torch.zeros(2, 64, 64) # Different batch size + + with pytest.raises(ValueError, match="same batch size"): + node.crop_to_mask_bounds(images, masks) + + +def test_crop_node_mismatched_spatial_dimensions(): + """Test error handling for mismatched spatial dimensions""" + from __init__ import CropImageBatchToMaskBoundsNodeAE + import torch + import pytest + + node = CropImageBatchToMaskBoundsNodeAE() + + # Mismatched spatial dimensions + images = torch.zeros(2, 64, 64, 3) + masks = torch.zeros(2, 32, 32) # Different spatial size + + with pytest.raises(ValueError, match="same spatial dimensions"): + node.crop_to_mask_bounds(images, masks) + + +def test_crop_node_invalid_tensor_types(): + """Test error handling for invalid tensor types""" + from __init__ import CropImageBatchToMaskBoundsNodeAE + import torch + import pytest + + node = CropImageBatchToMaskBoundsNodeAE() + + # Non-tensor images + with pytest.raises(ValueError, match="must be a torch.Tensor"): + node.crop_to_mask_bounds("not a tensor", torch.zeros(2, 64, 64)) + + # Non-tensor masks + with pytest.raises(ValueError, match="must be a torch.Tensor"): + node.crop_to_mask_bounds(torch.zeros(2, 64, 64, 3), "not a tensor") + + +def test_crop_node_invalid_tensor_shapes(): + """Test error handling for invalid tensor shapes""" + from __init__ import CropImageBatchToMaskBoundsNodeAE + import torch + import pytest + + node = CropImageBatchToMaskBoundsNodeAE() + + # Wrong image dimensions (3D instead of 4D) + with pytest.raises(ValueError, match="images must be 4D"): + node.crop_to_mask_bounds(torch.zeros(2, 64, 64), torch.zeros(2, 64, 64)) + + # Wrong mask dimensions (2D instead of 3D) + with pytest.raises(ValueError, match="masks must be 3D"): + node.crop_to_mask_bounds(torch.zeros(2, 64, 64, 3), torch.zeros(64, 64)) diff --git a/tests/test_image_selector.py b/tests/test_image_selector.py index f94cbcd..9fdd620 100644 --- a/tests/test_image_selector.py +++ b/tests/test_image_selector.py @@ -6,6 +6,9 @@ import os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import pytest +import torch +import numpy as np +from PIL import Image def test_module_imports(): @@ -49,3 +52,183 @@ def test_hex_to_rgb_invalid(): with pytest.raises(ValueError): hex_to_rgb("#FF") # Too short (2 chars) + + +def test_create_thumbnail(): + """Test thumbnail creation with various aspect ratios""" + from image_selector import create_thumbnail + + # Normal case - square image + img = Image.new('RGB', (100, 100), color='red') + result = create_thumbnail(img, (50, 50)) + assert result.size == (50, 50) + + # Wide image + img = Image.new('RGB', (200, 100), color='blue') + result = create_thumbnail(img, (50, 50)) + assert result.size[0] == 50 + assert result.size[1] <= 50 + + # Tall image + img = Image.new('RGB', (100, 200), color='green') + result = create_thumbnail(img, (50, 50)) + assert result.size[0] <= 50 + assert result.size[1] == 50 + + +def test_create_thumbnail_extreme_aspect_ratios(): + """Test thumbnail creation with extreme aspect ratios that could produce 0-pixel dimensions""" + from image_selector import create_thumbnail + + # Very wide image (1000x1) + img = Image.new('RGB', (1000, 1), color='blue') + result = create_thumbnail(img, (150, 150)) + assert result.size[0] == 150 + assert result.size[1] >= 1 # Should be at least 1 pixel + + # Very tall image (1x1000) + img = Image.new('RGB', (1, 1000), color='green') + result = create_thumbnail(img, (150, 150)) + assert result.size[0] >= 1 # Should be at least 1 pixel + assert result.size[1] == 150 + + # 1x1 pixel image + img = Image.new('RGB', (1, 1), color='red') + result = create_thumbnail(img, (150, 150)) + assert result.size[0] >= 1 + assert result.size[1] >= 1 + + +def test_pil_to_tensor(): + """Test PIL to tensor conversion""" + from image_selector import pil_to_tensor + + # RGB image + img = Image.new('RGB', (10, 10), color='red') + tensor, mask = pil_to_tensor(img) + assert tensor.shape == (1, 10, 10, 3) + assert mask.shape == (1, 10, 10) + assert torch.all(mask == 1.0) # All-white mask for RGB + + # RGBA image with transparency + img = Image.new('RGBA', (10, 10), color=(255, 0, 0, 128)) + tensor, mask = pil_to_tensor(img) + assert tensor.shape == (1, 10, 10, 3) + assert mask.shape == (1, 10, 10) + assert torch.all(mask <= 1.0) + assert torch.all(mask >= 0.0) + + +def test_pil_to_tensor_invalid(): + """Test pil_to_tensor with invalid inputs""" + from image_selector import pil_to_tensor + + # None input + with pytest.raises(ValueError, match="Image cannot be None"): + pil_to_tensor(None) + + # Non-PIL object + with pytest.raises(ValueError, match="Expected PIL Image"): + pil_to_tensor("not an image") + + with pytest.raises(ValueError, match="Expected PIL Image"): + pil_to_tensor(123) + + +def test_load_images_from_folder(): + """Test loading images from folder""" + from image_selector import load_images_from_folder + import tempfile + + # Empty folder + with tempfile.TemporaryDirectory() as tmpdir: + images, paths = load_images_from_folder(tmpdir) + assert len(images) == 0 + assert len(paths) == 0 + + # Non-existent folder + images, paths = load_images_from_folder('/nonexistent/path') + assert len(images) == 0 + assert len(paths) == 0 + + # Empty string + images, paths = load_images_from_folder('') + assert len(images) == 0 + assert len(paths) == 0 + + +def test_sort_image_paths(): + """Test sorting image paths""" + from image_selector import sort_image_paths + + # Empty list + result = sort_image_paths([], 'name', True) + assert result == [] + + # Single item + result = sort_image_paths(['/path/to/image.jpg'], 'name', True) + assert result == ['/path/to/image.jpg'] + + # Name sorting + paths = ['/path/c.jpg', '/path/a.jpg', '/path/b.jpg'] + result = sort_image_paths(paths, 'name', True) + assert result == ['/path/a.jpg', '/path/b.jpg', '/path/c.jpg'] + + # Descending + result = sort_image_paths(paths, 'name', False) + assert result == ['/path/c.jpg', '/path/b.jpg', '/path/a.jpg'] + + +def test_build_folder_tree(): + """Test folder tree building""" + from image_selector import build_folder_tree + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + # Create a simple folder structure + os.makedirs(os.path.join(tmpdir, 'subfolder1'), exist_ok=True) + os.makedirs(os.path.join(tmpdir, 'subfolder2'), exist_ok=True) + + tree = build_folder_tree(tmpdir, tmpdir) + + # Should have at least the parent and current folder + assert len(tree) >= 2 + + # First item should be parent navigation + assert tree[0][0].startswith('⬆️') + + # Second item should be current folder + assert tree[1][0].startswith('📁') + + # Non-existent path + tree = build_folder_tree('/nonexistent/path', '/nonexistent/path') + assert len(tree) >= 2 # Should still have navigation items + + +def test_process_selected_images_empty(): + """Test process_selected_images with empty input""" + from image_selector import process_selected_images + + images, masks, image_batch, mask_batch = process_selected_images([], [], None, None, '#000000') + + assert len(images) == 0 + assert len(masks) == 0 + assert image_batch.shape == (0, 0, 0, 3) + assert mask_batch.shape == (0, 0, 0) + + +def test_lazy_image_loader(): + """Test LazyImageLoader class""" + from image_selector import LazyImageLoader + + # Empty loader + loader = LazyImageLoader([]) + assert loader.paths == [] + + # Non-existent image + loader = LazyImageLoader(['/nonexistent.jpg']) + thumb = loader.get_thumbnail('/nonexistent.jpg') + assert thumb is not None # Should return error placeholder + + # Preload empty + loader.preload_visible([]) # Should not crash diff --git a/web/index.html b/web/index.html index 3c0c642..086daa3 100644 --- a/web/index.html +++ b/web/index.html @@ -84,16 +84,69 @@
-

🚀 Coming Soon

-

More interactive nodes for enhanced ComfyUI workflows

+

🖼️ Image Selector (Interactive)

+

Browse and select images from folders with an intuitive UI

+ +
+
📁 Folder navigation tree
+
⚡ Async loading for large folders
+
🔍 Multiple sort modes
+
👆 Click to select/deselect
+
🎨 Transparency support
+
📦 Batch tensor output
+
+ +

🎮 Controls

+
+

🚀 Z-Image Images To LoRA

+

Convert image batches to LoRA using DiffSynth-Studio

+ +
+
🎯 Single-node LoRA creation
+
🧠 Memory-efficient processing
+
📦 Batch processing support
+
⚖️ Weighted image averaging
+
🔢 Auto-version numbering
+
🔗 Direct integration with LoraLoader
+
+
+ +
+

✂️ Crop Image Batch to Mask Bounds

+

Intelligent mask-based cropping with consistent output dimensions

+ +
+
🎯 Auto-crop to mask boundaries
+
🔄 Individual or combined cropping
+
📐 Consistent output dimensions
+
🎨 Custom padding colors
+
🔄 Perfect mask alignment
+
📏 Optional forced dimensions
+
+
+ +
+

📊 Interpolate Float List

+

Process and interpolate comma-separated float lists

+ +
+
🔢 Linear interpolation
+
✂️ Trim or pad to target count
+
📋 CSV string output
+
🔄 Flexible formatting support
+
+
+

Made with ❤️ by by-ae

GitHub Repository