diff --git a/gates/loader.py b/gates/loader.py new file mode 100644 index 0000000..9ed70aa --- /dev/null +++ b/gates/loader.py @@ -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"} diff --git a/tests/test_loader.py b/tests/test_loader.py new file mode 100644 index 0000000..91f4f69 --- /dev/null +++ b/tests/test_loader.py @@ -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