feat: FolderImageLoader node (image/text/mask/filename/index)
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,70 @@
|
||||
# gates/loader.py
|
||||
import hashlib
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
from . import scan
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
def load_image_and_mask(path):
|
||||
img = Image.open(path)
|
||||
img = ImageOps.exif_transpose(img)
|
||||
arr = np.array(img.convert("RGB"), dtype=np.float32) / 255.0
|
||||
image = torch.from_numpy(arr).unsqueeze(0) # [1,H,W,3]
|
||||
h, w = arr.shape[0], arr.shape[1]
|
||||
if "A" in img.getbands():
|
||||
a = np.array(img.getchannel("A"), dtype=np.float32) / 255.0
|
||||
mask = (1.0 - torch.from_numpy(a)).unsqueeze(0) # [1,H,W]
|
||||
else:
|
||||
mask = torch.zeros((1, h, w), dtype=torch.float32)
|
||||
return image, mask
|
||||
|
||||
|
||||
class FolderImageLoader:
|
||||
CATEGORY = "Datasete Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "STRING", "MASK", "STRING", "INT")
|
||||
RETURN_NAMES = ("image", "text", "mask", "filename", "index")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"folder": ("STRING", {"default": ""}),
|
||||
"index": ("INT", {"default": 0, "min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
"control_after_generate": True}),
|
||||
"depth": ("INT", {"default": 0, "min": -1, "max": 64}),
|
||||
}
|
||||
}
|
||||
|
||||
def run(self, folder, index, depth=0):
|
||||
files = scan.list_images(folder, depth)
|
||||
idx = scan.resolve_index(len(files), index)
|
||||
path = files[idx]
|
||||
image, mask = load_image_and_mask(path)
|
||||
return (image, scan.read_sidecar(path), mask, scan.stem(path), idx)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, folder, index, depth=0, **kwargs):
|
||||
try:
|
||||
files = scan.list_images(folder, depth)
|
||||
idx = scan.resolve_index(len(files), index)
|
||||
path = files[idx]
|
||||
sc = scan.sidecar_path(path)
|
||||
parts = [folder, str(depth), str(idx),
|
||||
str(os.path.getmtime(path)),
|
||||
str(os.path.getmtime(sc)) if os.path.isfile(sc) else "0"]
|
||||
except Exception as e: # surface errors as a changed hash, not a crash here
|
||||
parts = [folder, str(depth), str(index), f"err:{e}"]
|
||||
return hashlib.sha256("|".join(parts).encode()).hexdigest()
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"FolderImageLoader": FolderImageLoader}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"FolderImageLoader": "Folder Image Loader"}
|
||||
@@ -0,0 +1,45 @@
|
||||
# tests/test_loader.py
|
||||
import io, os, torch
|
||||
from PIL import Image
|
||||
from gates import loader
|
||||
|
||||
def _save(path, color=(255, 0, 0), size=(4, 6), mode="RGB"): # size=(w,h)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
Image.new(mode, size, color).save(path)
|
||||
|
||||
def test_run_loads_image_text_stem_index(tmp_path):
|
||||
_save(str(tmp_path / "img1.png"), (255, 0, 0))
|
||||
_save(str(tmp_path / "img2.png"), (0, 255, 0))
|
||||
(tmp_path / "img2.txt").write_text("green frame\n", encoding="utf-8")
|
||||
n = loader.FolderImageLoader()
|
||||
image, text, mask, filename, index = n.run(folder=str(tmp_path), index=1, depth=0)
|
||||
assert image.shape == (1, 6, 4, 3)
|
||||
assert float(image[0, 0, 0, 1]) > 0.99 # green
|
||||
assert text == "green frame"
|
||||
assert filename == "img2"
|
||||
assert index == 1
|
||||
assert mask.shape == (1, 6, 4) and float(mask.max()) == 0.0 # no alpha -> zeros
|
||||
|
||||
def test_run_alpha_becomes_mask(tmp_path):
|
||||
# RGBA image, fully opaque alpha=255 -> mask = 1-1 = 0
|
||||
_save(str(tmp_path / "a.png"), (255, 255, 255, 255), mode="RGBA")
|
||||
n = loader.FolderImageLoader()
|
||||
_, _, mask, _, _ = n.run(folder=str(tmp_path), index=0, depth=0)
|
||||
assert float(mask.max()) == 0.0
|
||||
# transparent alpha=0 -> mask = 1-0 = 1
|
||||
_save(str(tmp_path / "b.png"), (255, 255, 255, 0), mode="RGBA")
|
||||
_, _, mask2, _, _ = n.run(folder=str(tmp_path), index=1, depth=0)
|
||||
assert float(mask2.min()) > 0.99
|
||||
|
||||
def test_run_out_of_range_raises(tmp_path):
|
||||
import pytest
|
||||
_save(str(tmp_path / "only.png"))
|
||||
n = loader.FolderImageLoader()
|
||||
with pytest.raises(IndexError):
|
||||
n.run(folder=str(tmp_path), index=9, depth=0)
|
||||
|
||||
def test_is_changed_differs_by_index_and_sidecar(tmp_path):
|
||||
_save(str(tmp_path / "img1.png")); _save(str(tmp_path / "img2.png"))
|
||||
h0 = loader.FolderImageLoader.IS_CHANGED(folder=str(tmp_path), index=0, depth=0)
|
||||
h1 = loader.FolderImageLoader.IS_CHANGED(folder=str(tmp_path), index=1, depth=0)
|
||||
assert h0 != h1
|
||||
Reference in New Issue
Block a user