feat: GridImagePool node (image/mask/index/count/label + IS_CHANGED)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-21 13:01:39 +02:00
parent 40be11cd95
commit 1db94dd57d
3 changed files with 127 additions and 1 deletions
+11
View File
@@ -0,0 +1,11 @@
"""Isolates the ComfyUI dependency so node.py stays unit-testable.
node.py imports ``grid_pool_base`` from here; tests monkeypatch
``node._grid_pool_base`` so ``folder_paths`` is never needed.
"""
import os
def grid_pool_base():
import folder_paths # imported lazily; only available inside ComfyUI
return os.path.join(folder_paths.get_input_directory(), "grid_pool")
+60 -1
View File
@@ -1,3 +1,62 @@
# gates/node.py — stub (filled in Task 9) """GridImagePool — the Image Pool (Grid) ComfyUI node."""
import os
from .gates_compat import grid_pool_base as _grid_pool_base
from . import pool, imaging
NODE_CLASS_MAPPINGS = {} NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {}
class GridImagePool:
CATEGORY = "Datasete Gates"
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
RETURN_NAMES = ("image", "mask", "index", "count", "label")
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"index": ("INT", {"default": -1, "min": -1, "max": 9999}),
},
"hidden": {"pool_id": "POOL_ID"},
}
@staticmethod
def _resolve(index, pool_id):
base = _grid_pool_base()
m = pool.read_manifest(base, pool_id)
idx = pool.resolve_slot(m, index)
return base, m, idx
def run(self, index, pool_id="default"):
base, m, idx = self._resolve(index, pool_id)
if idx < 0:
img, mask = imaging.empty_outputs()
return (img, mask, 0, 0, "")
slot = m["slots"][idx]
d = pool.pool_dir(base, pool_id)
img = imaging.load_image_tensor(str(d / slot["image"]))
h, w = int(img.shape[1]), int(img.shape[2])
mask_name = slot.get("mask")
mask = imaging.load_mask_tensor(str(d / mask_name) if mask_name else None, h, w)
return (img, mask, idx, len(m["slots"]), slot.get("label", ""))
@classmethod
def IS_CHANGED(cls, index, pool_id="default", **kwargs):
base, m, idx = cls._resolve(index, pool_id)
if idx < 0:
return imaging.change_hash(pool_id, -1, [])
slot = m["slots"][idx]
d = pool.pool_dir(base, pool_id)
mtimes = []
for key in ("image", "mask"):
name = slot.get(key)
p = d / name if name else None
mtimes.append(os.path.getmtime(p) if p and p.exists() else 0.0)
# include active so manual selection changes invalidate cache
return imaging.change_hash(pool_id, f"{idx}:{m.get('active')}", mtimes)
NODE_CLASS_MAPPINGS = {"GridImagePool": GridImagePool}
NODE_DISPLAY_NAME_MAPPINGS = {"GridImagePool": "Image Pool (Grid)"}
+56
View File
@@ -0,0 +1,56 @@
import io
import numpy as np, torch
from PIL import Image
from gates import node, pool
def _seed_pool(tmp_path, monkeypatch):
base = str(tmp_path / "grid_pool")
monkeypatch.setattr(node, "_grid_pool_base", lambda: base)
return base
def _add_png(base, pid, name_bytes_color, ts):
# write a real PNG via pool.add_image
buf = io.BytesIO(); Image.new("RGB", (4, 6), name_bytes_color).save(buf, "PNG")
return pool.add_image(base, pid, buf.getvalue(), ts=ts)
def test_execute_empty_pool_returns_blank(tmp_path, monkeypatch):
_seed_pool(tmp_path, monkeypatch)
n = node.GridImagePool()
img, mask, idx, count, label = n.run(index=-1, pool_id="p1")
assert img.shape == (1, 1, 1, 3)
assert count == 0 and idx == 0 and label == ""
def test_execute_selects_active(tmp_path, monkeypatch):
base = _seed_pool(tmp_path, monkeypatch)
_add_png(base, "p1", (255, 0, 0), 1)
_add_png(base, "p1", (0, 255, 0), 2)
pool.set_active(base, "p1", 1)
pool.set_label(base, "p1", 1, "green")
n = node.GridImagePool()
img, mask, idx, count, label = n.run(index=-1, pool_id="p1")
assert img.shape == (1, 6, 4, 3)
assert idx == 1 and count == 2 and label == "green"
assert float(img[0, 0, 0, 1]) > 0.99 # green channel
assert float(mask.max()) == 0.0 # no mask yet
def test_execute_forced_index_clamps(tmp_path, monkeypatch):
base = _seed_pool(tmp_path, monkeypatch)
_add_png(base, "p1", (255, 0, 0), 1)
n = node.GridImagePool()
_, _, idx, count, _ = n.run(index=9, pool_id="p1")
assert idx == 0 and count == 1
def test_is_changed_differs_after_active_change(tmp_path, monkeypatch):
base = _seed_pool(tmp_path, monkeypatch)
_add_png(base, "p1", (255, 0, 0), 1)
_add_png(base, "p1", (0, 255, 0), 2)
h1 = node.GridImagePool.IS_CHANGED(index=-1, pool_id="p1")
pool.set_active(base, "p1", 1)
h2 = node.GridImagePool.IS_CHANGED(index=-1, pool_id="p1")
assert h1 != h2