127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
"""
|
|
Example: Basic Image Processor Node
|
|
|
|
A complete working example of a ComfyUI node that resizes images
|
|
with various interpolation methods.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
class ExampleImageResizer:
|
|
"""
|
|
Resize images to target dimensions with multiple interpolation options.
|
|
|
|
This demonstrates:
|
|
- Basic node structure
|
|
- IMAGE input/output types
|
|
- Dropdown widget configuration
|
|
- PIL-based image processing
|
|
- Batch processing support
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {
|
|
"tooltip": "Input image tensor (B, H, W, C)"
|
|
}),
|
|
"width": ("INT", {
|
|
"default": 512,
|
|
"min": 64,
|
|
"max": 8192,
|
|
"step": 64,
|
|
"tooltip": "Target width in pixels"
|
|
}),
|
|
"height": ("INT", {
|
|
"default": 512,
|
|
"min": 64,
|
|
"max": 8192,
|
|
"step": 64,
|
|
"tooltip": "Target height in pixels"
|
|
}),
|
|
"interpolation": (["nearest", "bilinear", "bicubic", "lanczos"], {
|
|
"default": "bilinear",
|
|
"tooltip": "Resampling method"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"maintain_aspect": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Keep original aspect ratio"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("resized_image",)
|
|
FUNCTION = "resize"
|
|
CATEGORY = "example/image"
|
|
|
|
# Maps string names to PIL constants
|
|
INTERPOLATION_MAP = {
|
|
"nearest": Image.NEAREST,
|
|
"bilinear": Image.BILINEAR,
|
|
"bicubic": Image.BICUBIC,
|
|
"lanczos": Image.LANCZOS,
|
|
}
|
|
|
|
def resize(self, image, width, height, interpolation, maintain_aspect=False):
|
|
"""
|
|
Resize input image(s) to target dimensions.
|
|
|
|
Args:
|
|
image: Tensor of shape (B, H, W, C) with values in [0, 1]
|
|
width: Target width
|
|
height: Target height
|
|
interpolation: Resampling method name
|
|
maintain_aspect: Whether to preserve aspect ratio
|
|
|
|
Returns:
|
|
Tuple containing resized image tensor
|
|
"""
|
|
# Get interpolation method
|
|
interp_method = self.INTERPOLATION_MAP.get(interpolation, Image.BILINEAR)
|
|
|
|
batch_size, orig_h, orig_w, channels = image.shape
|
|
|
|
# Calculate dimensions if maintaining aspect ratio
|
|
if maintain_aspect:
|
|
aspect = orig_w / orig_h
|
|
if width / height > aspect:
|
|
width = int(height * aspect)
|
|
else:
|
|
height = int(width / aspect)
|
|
|
|
# Process each image in batch
|
|
resized_images = []
|
|
for i in range(batch_size):
|
|
# Convert tensor to PIL (0-255 range)
|
|
img_np = (image[i].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(img_np)
|
|
|
|
# Resize
|
|
pil_img = pil_img.resize((width, height), interp_method)
|
|
|
|
# Convert back to tensor (0-1 range)
|
|
img_np = np.array(pil_img).astype(np.float32) / 255.0
|
|
resized_images.append(torch.from_numpy(img_np))
|
|
|
|
# Stack back into batch
|
|
result = torch.stack(resized_images)
|
|
|
|
return (result,)
|
|
|
|
|
|
# Node registration (would go in __init__.py)
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ExampleImageResizer": ExampleImageResizer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ExampleImageResizer": "Example: Image Resizer",
|
|
}
|