341 lines
14 KiB
Python
341 lines
14 KiB
Python
"""
|
||
ComfyUI-SmartGridCutter — Universal Grid Image Splitter
|
||
|
||
Automatically detects grid layout by analyzing contrast/brightness
|
||
and splits any grid image into individual frames.
|
||
|
||
Works with any grid: 3×3 storyboards, 4×2 character sheets,
|
||
2×3 panels, or any arbitrary grid layout.
|
||
|
||
Author: iMind Studio
|
||
License: MIT
|
||
"""
|
||
|
||
import torch
|
||
import numpy as np
|
||
from PIL import Image
|
||
import io
|
||
import urllib.request
|
||
|
||
|
||
class SmartGridCutter:
|
||
"""
|
||
Splits a grid image into individual cells by detecting
|
||
divider lines through brightness/contrast analysis.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"mode": (["auto", "3x3", "4x2", "3x2", "2x2", "custom"],),
|
||
},
|
||
"optional": {
|
||
"custom_cols": ("INT", {"default": 3, "min": 1, "max": 10, "step": 1}),
|
||
"custom_rows": ("INT", {"default": 3, "min": 1, "max": 10, "step": 1}),
|
||
"divider_threshold": ("FLOAT", {"default": 0.15, "min": 0.01, "max": 0.5, "step": 0.01,
|
||
"tooltip": "Brightness difference threshold for detecting divider lines. Lower = more sensitive."}),
|
||
"min_divider_width": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1,
|
||
"tooltip": "Minimum width in pixels for a line to be considered a divider."}),
|
||
"padding": ("INT", {"default": 2, "min": 0, "max": 20, "step": 1,
|
||
"tooltip": "Pixels to trim from each cell edge after cutting."}),
|
||
"align_to": ("INT", {"default": 16, "min": 1, "max": 64, "step": 1,
|
||
"tooltip": "Align cell dimensions to this value (16 for video compatibility)."}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||
"IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||
"INT", "INT", "INT", "STRING")
|
||
RETURN_NAMES = ("frame_0", "frame_1", "frame_2", "frame_3", "frame_4",
|
||
"frame_5", "frame_6", "frame_7", "frame_8",
|
||
"cols", "rows", "total_frames", "grid_info")
|
||
|
||
FUNCTION = "cut_grid"
|
||
CATEGORY = "image/transform"
|
||
DESCRIPTION = "Splits a grid image into individual frames by detecting divider lines. Works with any grid layout."
|
||
|
||
def _find_dividers(self, gray_array, axis, threshold, min_width, margin_pct=0.05):
|
||
"""
|
||
Find divider positions along an axis using variance + contrast analysis.
|
||
|
||
Dividers are detected by TWO criteria:
|
||
1. Low variance (uniform color — solid white/black/gray line)
|
||
2. Contrast with neighbors (different brightness from surrounding cells)
|
||
|
||
Args:
|
||
gray_array: 2D numpy array (grayscale, 0-1 float)
|
||
axis: 0 for horizontal dividers (scanning rows), 1 for vertical (scanning cols)
|
||
threshold: NOT USED directly anymore, kept for API compat
|
||
min_width: minimum pixel width for a valid divider
|
||
margin_pct: ignore edges (percentage of total size)
|
||
"""
|
||
h, w = gray_array.shape
|
||
|
||
if axis == 1: # vertical dividers (scan columns)
|
||
size = w
|
||
# Per-column stats
|
||
profile_mean = np.mean(gray_array[::2, :], axis=0)
|
||
profile_std = np.std(gray_array[::2, :], axis=0)
|
||
else: # horizontal dividers (scan rows)
|
||
size = h
|
||
profile_mean = np.mean(gray_array[:, ::2], axis=1)
|
||
profile_std = np.std(gray_array[:, ::2], axis=1)
|
||
|
||
# Dividers have LOW std (uniform color across the perpendicular axis)
|
||
# Content cells have HIGH std (varied pixel values)
|
||
median_std = np.median(profile_std)
|
||
# A divider line's std should be much lower than typical content
|
||
std_threshold = max(median_std * 0.3, 0.01)
|
||
|
||
margin = int(size * margin_pct)
|
||
dividers = []
|
||
in_divider = False
|
||
start = 0
|
||
|
||
for i in range(margin, size - margin):
|
||
is_uniform = profile_std[i] < std_threshold
|
||
|
||
if is_uniform:
|
||
if not in_divider:
|
||
start = i
|
||
in_divider = True
|
||
else:
|
||
if in_divider:
|
||
width = i - start
|
||
if width >= min_width:
|
||
# Verify it's actually a divider by checking contrast with neighbors
|
||
div_brightness = np.mean(profile_mean[start:i])
|
||
# Check brightness of cells on both sides
|
||
left_zone = max(0, start - min(30, start))
|
||
right_zone = min(size, i + 30)
|
||
neighbor_brightness = np.mean(
|
||
np.concatenate([
|
||
profile_mean[left_zone:start],
|
||
profile_mean[i:right_zone]
|
||
])
|
||
) if (start > left_zone or i < right_zone) else div_brightness
|
||
|
||
contrast = abs(div_brightness - neighbor_brightness)
|
||
if contrast > 0.05: # Must have SOME contrast with neighbors
|
||
dividers.append({
|
||
"start": start,
|
||
"end": i,
|
||
"mid": (start + i) // 2,
|
||
"width": width,
|
||
})
|
||
in_divider = False
|
||
|
||
return dividers
|
||
|
||
def _build_grid_from_dividers(self, x_dividers, y_dividers, img_w, img_h):
|
||
"""Build cell boundaries from detected dividers."""
|
||
cols = len(x_dividers) + 1
|
||
rows = len(y_dividers) + 1
|
||
|
||
# X boundaries: [0, end_of_div1, end_of_div2, ..., img_w]
|
||
x_starts = [0] + [d["end"] for d in x_dividers]
|
||
x_ends = [d["start"] for d in x_dividers] + [img_w]
|
||
|
||
y_starts = [0] + [d["end"] for d in y_dividers]
|
||
y_ends = [d["start"] for d in y_dividers] + [img_h]
|
||
|
||
return cols, rows, x_starts, x_ends, y_starts, y_ends
|
||
|
||
def _build_grid_from_preset(self, cols, rows, img_w, img_h):
|
||
"""Build cell boundaries using simple division."""
|
||
cell_w = img_w // cols
|
||
cell_h = img_h // rows
|
||
|
||
x_starts = [i * cell_w for i in range(cols)]
|
||
x_ends = [(i + 1) * cell_w for i in range(cols)]
|
||
y_starts = [i * cell_h for i in range(rows)]
|
||
y_ends = [(i + 1) * cell_h for i in range(rows)]
|
||
|
||
return x_starts, x_ends, y_starts, y_ends
|
||
|
||
def _tensor_to_pil(self, tensor):
|
||
"""Convert ComfyUI IMAGE tensor to PIL Image."""
|
||
# ComfyUI images: [B, H, W, C], float32, 0-1
|
||
if tensor.dim() == 4:
|
||
tensor = tensor[0] # Take first batch
|
||
arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||
return Image.fromarray(arr)
|
||
|
||
def _pil_to_tensor(self, pil_img):
|
||
"""Convert PIL Image to ComfyUI IMAGE tensor."""
|
||
arr = np.array(pil_img).astype(np.float32) / 255.0
|
||
if arr.ndim == 2:
|
||
arr = np.stack([arr] * 3, axis=-1)
|
||
return torch.from_numpy(arr).unsqueeze(0) # [1, H, W, C]
|
||
|
||
def cut_grid(self, image, mode, custom_cols=3, custom_rows=3,
|
||
divider_threshold=0.15, min_divider_width=2,
|
||
padding=2, align_to=16):
|
||
"""Main cutting function."""
|
||
|
||
pil_img = self._tensor_to_pil(image)
|
||
img_w, img_h = pil_img.size
|
||
|
||
# Convert to grayscale for analysis
|
||
gray = np.array(pil_img.convert("L")).astype(np.float32) / 255.0
|
||
|
||
# Determine grid layout
|
||
presets = {
|
||
"3x3": (3, 3),
|
||
"4x2": (4, 2),
|
||
"3x2": (3, 2),
|
||
"2x2": (2, 2),
|
||
"custom": (custom_cols, custom_rows),
|
||
}
|
||
|
||
if mode == "auto":
|
||
# Smart detection
|
||
x_dividers = self._find_dividers(gray, axis=1,
|
||
threshold=divider_threshold, min_width=min_divider_width)
|
||
y_dividers = self._find_dividers(gray, axis=0,
|
||
threshold=divider_threshold, min_width=min_divider_width)
|
||
|
||
cols, rows, x_starts, x_ends, y_starts, y_ends = \
|
||
self._build_grid_from_dividers(x_dividers, y_dividers, img_w, img_h)
|
||
|
||
detect_method = "auto_detected"
|
||
|
||
# Sanity check: if detection gives weird results (1×1 or 20×20), fallback
|
||
if cols < 2 and rows < 2:
|
||
# Try preset-based detection using aspect ratio
|
||
ar = img_w / img_h
|
||
if 0.9 <= ar <= 1.1:
|
||
cols, rows = 3, 3 # Square → likely 3×3
|
||
elif ar > 1.4:
|
||
cols, rows = 4, 2 # Wide → likely 4×2
|
||
else:
|
||
cols, rows = 3, 2 # Portrait-ish → likely 3×2
|
||
|
||
x_starts, x_ends, y_starts, y_ends = \
|
||
self._build_grid_from_preset(cols, rows, img_w, img_h)
|
||
detect_method = "aspect_ratio_fallback"
|
||
|
||
else:
|
||
cols, rows = presets[mode]
|
||
# Try smart detection first, validate against preset
|
||
x_dividers = self._find_dividers(gray, axis=1,
|
||
threshold=divider_threshold, min_width=min_divider_width)
|
||
y_dividers = self._find_dividers(gray, axis=0,
|
||
threshold=divider_threshold, min_width=min_divider_width)
|
||
|
||
detected_cols = len(x_dividers) + 1
|
||
detected_rows = len(y_dividers) + 1
|
||
|
||
if detected_cols == cols and detected_rows == rows:
|
||
# Smart detection matches preset — use smart bounds
|
||
_, _, x_starts, x_ends, y_starts, y_ends = \
|
||
self._build_grid_from_dividers(x_dividers, y_dividers, img_w, img_h)
|
||
detect_method = "smart_verified"
|
||
else:
|
||
# Fallback to simple division
|
||
x_starts, x_ends, y_starts, y_ends = \
|
||
self._build_grid_from_preset(cols, rows, img_w, img_h)
|
||
detect_method = f"preset_fallback (detected {detected_cols}x{detected_rows})"
|
||
|
||
# Cut frames
|
||
frames = []
|
||
for r in range(rows):
|
||
for c in range(cols):
|
||
left = x_starts[c] + padding
|
||
top = y_starts[r] + padding
|
||
right = x_ends[c] - padding
|
||
bottom = y_ends[r] - padding
|
||
|
||
w = right - left
|
||
h = bottom - top
|
||
|
||
# Align dimensions
|
||
if align_to > 1:
|
||
w = (w // align_to) * align_to
|
||
h = (h // align_to) * align_to
|
||
|
||
cell = pil_img.crop((left, top, left + w, top + h))
|
||
frames.append(self._pil_to_tensor(cell))
|
||
|
||
total = len(frames)
|
||
|
||
# Build info string
|
||
info = (f"Grid: {cols}×{rows} = {total} frames | "
|
||
f"Source: {img_w}×{img_h} | "
|
||
f"Cell: {frames[0].shape[2]}×{frames[0].shape[1]} | "
|
||
f"Detection: {detect_method}")
|
||
|
||
# Pad to 9 outputs (max supported)
|
||
while len(frames) < 9:
|
||
# Create a 1×1 black tensor as placeholder
|
||
frames.append(torch.zeros(1, 16, 16, 3))
|
||
|
||
return (*frames[:9], cols, rows, total, info)
|
||
|
||
|
||
class SmartGridCutterURL:
|
||
"""
|
||
Downloads an image from URL and splits it into grid cells.
|
||
Useful for n8n / API integration.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image_url": ("STRING", {"default": "", "multiline": False}),
|
||
"mode": (["auto", "3x3", "4x2", "3x2", "2x2", "custom"],),
|
||
},
|
||
"optional": {
|
||
"custom_cols": ("INT", {"default": 3, "min": 1, "max": 10}),
|
||
"custom_rows": ("INT", {"default": 3, "min": 1, "max": 10}),
|
||
"divider_threshold": ("FLOAT", {"default": 0.15, "min": 0.01, "max": 0.5, "step": 0.01}),
|
||
"padding": ("INT", {"default": 2, "min": 0, "max": 20}),
|
||
"align_to": ("INT", {"default": 16, "min": 1, "max": 64}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||
"IMAGE", "IMAGE", "IMAGE", "IMAGE",
|
||
"INT", "INT", "INT", "STRING")
|
||
RETURN_NAMES = ("frame_0", "frame_1", "frame_2", "frame_3", "frame_4",
|
||
"frame_5", "frame_6", "frame_7", "frame_8",
|
||
"cols", "rows", "total_frames", "grid_info")
|
||
|
||
FUNCTION = "cut_from_url"
|
||
CATEGORY = "image/transform"
|
||
DESCRIPTION = "Downloads image from URL and splits it into grid frames."
|
||
|
||
def cut_from_url(self, image_url, mode, custom_cols=3, custom_rows=3,
|
||
divider_threshold=0.15, padding=2, align_to=16):
|
||
"""Download image from URL and cut it."""
|
||
# Download image
|
||
req = urllib.request.Request(image_url, headers={"User-Agent": "ComfyUI-SmartGridCutter/1.0"})
|
||
with urllib.request.urlopen(req, timeout=30) as response:
|
||
img_data = response.read()
|
||
|
||
pil_img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||
|
||
# Convert to tensor
|
||
arr = np.array(pil_img).astype(np.float32) / 255.0
|
||
tensor = torch.from_numpy(arr).unsqueeze(0)
|
||
|
||
# Use the main cutter
|
||
cutter = SmartGridCutter()
|
||
return cutter.cut_grid(tensor, mode, custom_cols, custom_rows,
|
||
divider_threshold, 2, padding, align_to)
|
||
|
||
|
||
# ── Registration ─────────────────────────────────────
|
||
NODE_CLASS_MAPPINGS = {
|
||
"SmartGridCutter": SmartGridCutter,
|
||
"SmartGridCutterURL": SmartGridCutterURL,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"SmartGridCutter": "Smart Grid Cutter ✂️",
|
||
"SmartGridCutterURL": "Smart Grid Cutter (URL) ✂️🔗",
|
||
}
|