Merge pull request #6 from by-ae/cursor/quality-fixes-8b9c
Quality fixes: input validation and comprehensive tests
This commit is contained in:
cursor[bot]
+19
@@ -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}")
|
||||
|
||||
+21
-1
@@ -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')
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
+59
-6
@@ -84,16 +84,69 @@
|
||||
</div>
|
||||
|
||||
<div class="node-card">
|
||||
<h2>🚀 Coming Soon</h2>
|
||||
<p>More interactive nodes for enhanced ComfyUI workflows</p>
|
||||
<h2>🖼️ Image Selector (Interactive)</h2>
|
||||
<p>Browse and select images from folders with an intuitive UI</p>
|
||||
|
||||
<div class="feature-list">
|
||||
<div class="feature-item">📁 Folder navigation tree</div>
|
||||
<div class="feature-item">⚡ Async loading for large folders</div>
|
||||
<div class="feature-item">🔍 Multiple sort modes</div>
|
||||
<div class="feature-item">👆 Click to select/deselect</div>
|
||||
<div class="feature-item">🎨 Transparency support</div>
|
||||
<div class="feature-item">📦 Batch tensor output</div>
|
||||
</div>
|
||||
|
||||
<h3>🎮 Controls</h3>
|
||||
<ul>
|
||||
<li>Real-time video processing nodes</li>
|
||||
<li>Interactive mask editors</li>
|
||||
<li>Streaming workflow controllers</li>
|
||||
<li>And more...</li>
|
||||
<li><strong>Left Click:</strong> Select/deselect images or navigate folders</li>
|
||||
<li><strong>Ctrl + A:</strong> Select/deselect all images</li>
|
||||
<li><strong>Ctrl + S:</strong> Cycle sort modes</li>
|
||||
<li><strong>Ctrl + D:</strong> Toggle sort direction</li>
|
||||
<li><strong>Mouse Wheel:</strong> Scroll images or folders</li>
|
||||
<li><strong>ESC/ENTER:</strong> Finish selection</li>
|
||||
</ul>
|
||||
</div>
|
||||
|
||||
<div class="node-card">
|
||||
<h2>🚀 Z-Image Images To LoRA</h2>
|
||||
<p>Convert image batches to LoRA using DiffSynth-Studio</p>
|
||||
|
||||
<div class="feature-list">
|
||||
<div class="feature-item">🎯 Single-node LoRA creation</div>
|
||||
<div class="feature-item">🧠 Memory-efficient processing</div>
|
||||
<div class="feature-item">📦 Batch processing support</div>
|
||||
<div class="feature-item">⚖️ Weighted image averaging</div>
|
||||
<div class="feature-item">🔢 Auto-version numbering</div>
|
||||
<div class="feature-item">🔗 Direct integration with LoraLoader</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="node-card">
|
||||
<h2>✂️ Crop Image Batch to Mask Bounds</h2>
|
||||
<p>Intelligent mask-based cropping with consistent output dimensions</p>
|
||||
|
||||
<div class="feature-list">
|
||||
<div class="feature-item">🎯 Auto-crop to mask boundaries</div>
|
||||
<div class="feature-item">🔄 Individual or combined cropping</div>
|
||||
<div class="feature-item">📐 Consistent output dimensions</div>
|
||||
<div class="feature-item">🎨 Custom padding colors</div>
|
||||
<div class="feature-item">🔄 Perfect mask alignment</div>
|
||||
<div class="feature-item">📏 Optional forced dimensions</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="node-card">
|
||||
<h2>📊 Interpolate Float List</h2>
|
||||
<p>Process and interpolate comma-separated float lists</p>
|
||||
|
||||
<div class="feature-list">
|
||||
<div class="feature-item">🔢 Linear interpolation</div>
|
||||
<div class="feature-item">✂️ Trim or pad to target count</div>
|
||||
<div class="feature-item">📋 CSV string output</div>
|
||||
<div class="feature-item">🔄 Flexible formatting support</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style="text-align: center; margin-top: 40px; color: #666;">
|
||||
<p>Made with ❤️ by <a href="https://github.com/by-ae" style="color: #667eea;">by-ae</a></p>
|
||||
<p><a href="https://github.com/by-ae/ae-in-workflow">GitHub Repository</a></p>
|
||||
|
||||
Reference in New Issue
Block a user