Add file I/O support for strings, images, and masks

Extended functionality to load and save strings, RGB/RGBA images, and masks, with various format options and error handling. Updated tests to cover new capabilities, ensuring robust behavior across different use cases. Updated documentation and version to reflect these changes.

README.md

@@ -173,6 +173,16 @@ These nodes are very lightweight and require no additional dependencies.
 - Path conversions:
     - relative - Computes a relative path from a start path to a target path
     - expand_vars - Replaces environment variables in a path with their values
+- File loading:
+    - load STRING from file - Loads a text file and returns its content as a STRING
+    - load IMAGE from file (RGB) - Loads an image and returns RGB channels as a tensor
+    - load IMAGE+MASK from file (RGBA) - Loads an image and returns RGB channels as a tensor and alpha channel as a mask
+    - load MASK from alpha channel - Loads an image and extracts its alpha channel as a mask
+    - load MASK from greyscale/red - Loads an image and creates a mask from its greyscale or red channel
+- File saving:
+    - save STRING to file - Saves a string to a text file with optional directory creation
+    - save IMAGE to file - Saves an image tensor to a file in various formats (PNG, JPG, WEBP, JXL)
+    - save IMAGE+MASK to file - Saves an image with transparency using a mask as the alpha channel
 
 ### SET: Python set manipulation nodes (as a single variable)
 - Creation:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "basic_data_handling"
-version = "0.4.3"
+version = "0.4.4"
 description = """Basic Python functions for manipulating data that every programmer is used to, lightweight with no additional dependencies.
 
 Supported data types:
@@ -20,7 +20,7 @@ Feature categories:
 - Mathematical operations
 - Mathematical formula node in a safe implementation
 - String manipulation
-- File system path handling
+- File system path handling, including STRING, IMAGE and MASK load and save
 - SET operations"""
 authors = [
   {name = "StableLlama"}

src/basic_data_handling/path_nodes.py

@@ -11,6 +11,8 @@ except:
         FLOAT = "FLOAT"
         STRING = "STRING"
         NUMBER = "FLOAT,INT"
+        IMAGE = "IMAGE"
+        MASK = "MASK"
         ANY = "*"
     ComfyNodeABC = object
 

tests/test_path_nodes.py

@@ -1,12 +1,17 @@
 import os
 import pytest
 import platform
+import numpy as np
+import torch
+from PIL import Image
 
 from src.basic_data_handling.path_nodes import (
     PathJoin, PathAbspath, PathExists, PathIsFile, PathIsDir, PathGetSize,
     PathSplit, PathSplitExt, PathBasename, PathDirname, PathGetExtension,
     PathSetExtension, PathNormalize, PathRelative, PathGlob, PathExpandVars, PathGetCwd,
-    PathListDir, PathIsAbsolute, PathCommonPrefix,
+    PathListDir, PathIsAbsolute, PathCommonPrefix, PathLoadStringFile, PathSaveStringFile,
+    PathLoadImageRGB, PathSaveImageRGB, PathLoadImageRGBA, PathSaveImageRGBA,
+    PathLoadMaskFromAlpha, PathLoadMaskFromGreyscale,
 )
 
 
@@ -21,6 +26,39 @@ def test_path_join():
     nested_path = os.path.join("folder", "subfolder", "file.txt")
     assert node.join_paths("folder", os.path.join("subfolder", "file.txt")) == (nested_path,)
 
+def test_path_load_save_string_file(tmp_path):
+    # Test saving a string to a file
+    save_node = PathSaveStringFile()
+    test_string = "This is a test string\nwith multiple lines"
+    file_path = str(tmp_path / "test_string.txt")
+
+    # Save the string
+    assert save_node.save_text(test_string, file_path) == (True,)
+
+    # Verify file exists
+    assert os.path.exists(file_path)
+
+    # Test loading the string back
+    load_node = PathLoadStringFile()
+    loaded_string = load_node.load_text(file_path)
+    assert loaded_string == (test_string,)
+
+    # Test creating directories when saving
+    nested_path = str(tmp_path / "nested" / "dir" / "test.txt")
+    assert save_node.save_text("Test with nested dirs", nested_path) == (True,)
+    assert os.path.exists(nested_path)
+
+    # Test different encodings
+    utf8_text = "UTF-8 text with special chars: 你好, ñ, é, ö"
+    utf8_path = str(tmp_path / "utf8_test.txt")
+    assert save_node.save_text(utf8_text, utf8_path, encoding="utf-8") == (True,)
+    assert load_node.load_text(utf8_path) == (utf8_text,)
+
+    # Test error handling
+    with pytest.raises(FileNotFoundError):
+        load_node.load_text(str(tmp_path / "nonexistent.txt"))
+
+
 
 def test_path_abspath():
     node = PathAbspath()
@@ -217,6 +255,193 @@ def test_path_set_extension():
     assert node.set_extension("", ".txt") == (".txt",)
 
 
+def test_path_load_mask_nodes(tmp_path, monkeypatch):
+    # Create test images
+    img_size = (64, 64)
+
+    # Image with alpha channel
+    rgba_img = Image.new('RGBA', img_size, color=(255, 0, 0, 128))  # Semi-transparent red
+    rgba_path = str(tmp_path / "alpha_mask.png")
+    rgba_img.save(rgba_path)
+
+    # Grayscale image
+    gray_img = Image.new('L', img_size)
+    # Create a gradient from black to white
+    for y in range(img_size[1]):
+        for x in range(img_size[0]):
+            gray_img.putpixel((x, y), int(255 * (x / img_size[0])))
+    gray_path = str(tmp_path / "gray_mask.png")
+    gray_img.save(gray_path)
+
+    # Mock the helper functions
+    def mock_load_image_helper(path):
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Image file not found: {path}")
+        return Image.open(path)
+
+    def mock_extract_mask_from_alpha(img):
+        if 'A' in img.getbands():
+            alpha = np.array(img.getchannel('A')).astype(np.float32) / 255.0
+            mask_tensor = 1.0 - torch.from_numpy(alpha)
+            return mask_tensor.unsqueeze(0)
+        return torch.zeros((1, img.height, img.width), dtype=torch.float32)
+
+    def mock_extract_mask_from_greyscale(img):
+        if img.mode == 'L':
+            gray = np.array(img).astype(np.float32) / 255.0
+        else:
+            gray = np.array(img.getchannel('R')).astype(np.float32) / 255.0
+        mask_tensor = 1.0 - torch.from_numpy(gray)
+        return mask_tensor.unsqueeze(0)
+
+    monkeypatch.setattr("src.basic_data_handling.path_nodes.load_image_helper", mock_load_image_helper)
+    monkeypatch.setattr("src.basic_data_handling.path_nodes.extract_mask_from_alpha", mock_extract_mask_from_alpha)
+    monkeypatch.setattr("src.basic_data_handling.path_nodes.extract_mask_from_greyscale", mock_extract_mask_from_greyscale)
+
+    # Test loading mask from alpha channel
+    alpha_node = PathLoadMaskFromAlpha()
+    alpha_mask = alpha_node.load_mask_from_alpha(rgba_path)
+
+    # Verify the mask shape
+    assert isinstance(alpha_mask, tuple)
+    assert len(alpha_mask) == 1
+    assert isinstance(alpha_mask[0], torch.Tensor)
+    assert alpha_mask[0].shape == (1, img_size[1], img_size[0])
+
+    # Test loading mask from grayscale image
+    gray_node = PathLoadMaskFromGreyscale()
+    gray_mask = gray_node.load_mask_from_greyscale(gray_path)
+
+    # Verify the mask shape
+    assert isinstance(gray_mask, tuple)
+    assert len(gray_mask) == 1
+    assert isinstance(gray_mask[0], torch.Tensor)
+    assert gray_mask[0].shape == (1, img_size[1], img_size[0])
+
+    # Test with invert option
+    inverted_mask = gray_node.load_mask_from_greyscale(gray_path, invert=True)
+    assert torch.allclose(inverted_mask[0], 1.0 - gray_mask[0])
+
+    # Test error handling
+    with pytest.raises(FileNotFoundError):
+        alpha_node.load_mask_from_alpha(str(tmp_path / "nonexistent.png"))
+
+
+def test_path_load_save_image_rgba(tmp_path, monkeypatch):
+    # Create a test image with transparency
+    img_size = (64, 64)
+    test_img = Image.new('RGBA', img_size, color=(255, 0, 0, 128))  # Semi-transparent red
+    img_path = str(tmp_path / "test_rgba.png")
+    test_img.save(img_path)
+
+    # Mock the load_image_helper and extraction functions
+    def mock_load_image_helper(path):
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Image file not found: {path}")
+        return Image.open(path)
+
+    def mock_extract_mask_from_alpha(img):
+        if 'A' in img.getbands():
+            alpha = np.array(img.getchannel('A')).astype(np.float32) / 255.0
+            mask_tensor = 1.0 - torch.from_numpy(alpha)
+            return mask_tensor.unsqueeze(0)
+        return torch.zeros((1, img.height, img.width), dtype=torch.float32)
+
+    monkeypatch.setattr("src.basic_data_handling.path_nodes.load_image_helper", mock_load_image_helper)
+    monkeypatch.setattr("src.basic_data_handling.path_nodes.extract_mask_from_alpha", mock_extract_mask_from_alpha)
+
+    # Test loading an image with alpha
+    load_node = PathLoadImageRGBA()
+    loaded_img, loaded_mask = load_node.load_image_rgba(img_path)
+
+    # Verify that the returned objects are tensors with the right shapes
+    assert isinstance(loaded_img, torch.Tensor)
+    assert isinstance(loaded_mask, torch.Tensor)
+    assert loaded_img.shape == (1, img_size[1], img_size[0], 3)  # (batch, height, width, channels)
+    assert loaded_mask.shape == (1, img_size[1], img_size[0])    # (batch, height, width)
+
+    # Test saving an image with mask
+    save_node = PathSaveImageRGBA()
+    output_path = str(tmp_path / "output_rgba")
+
+    # Create a simple red test image tensor and a gradient mask
+    red_img = torch.zeros(1, img_size[1], img_size[0], 3)
+    red_img[0, :, :, 0] = 1.0  # Red channel set to 1
+
+    # Create a gradient mask (0.0 to 1.0 from left to right)
+    mask = torch.zeros(1, img_size[1], img_size[0])
+    for i in range(img_size[0]):
+        mask[0, :, i] = i / img_size[0]
+
+    # Save the image with mask
+    assert save_node.save_image_with_mask(red_img, mask, output_path) == (True,)
+
+    # Verify the image was saved with the correct extension
+    assert os.path.exists(output_path + ".png")
+
+    # Test with invert_mask option
+    assert save_node.save_image_with_mask(red_img, mask, str(tmp_path / "inverted"), invert_mask=True) == (True,)
+    assert os.path.exists(str(tmp_path / "inverted.png"))
+
+    # Test with JPEG format (should switch to PNG for transparency)
+    assert save_node.save_image_with_mask(red_img, mask, str(tmp_path / "jpeg_test"), format="jpg") == (True,)
+    # Should be saved as PNG despite the request for JPEG
+    assert os.path.exists(str(tmp_path / "jpeg_test.png"))
+
+
+def test_path_load_save_image_rgb(tmp_path, monkeypatch):
+    # Create a test image
+    img_size = (64, 64)
+    test_img = Image.new('RGB', img_size, color='red')
+    img_path = str(tmp_path / "test_rgb.png")
+    test_img.save(img_path)
+
+    # Mock the load_image_helper function to avoid PIL import issues in testing
+    def mock_load_image_helper(path):
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Image file not found: {path}")
+        return Image.open(path)
+
+    monkeypatch.setattr("src.basic_data_handling.path_nodes.load_image_helper", mock_load_image_helper)
+
+    # Test loading an image
+    load_node = PathLoadImageRGB()
+    loaded_img = load_node.load_image_rgb(img_path)
+
+    # Verify that the returned object is a tensor with the right shape
+    assert isinstance(loaded_img, tuple)
+    assert len(loaded_img) == 1
+    assert isinstance(loaded_img[0], torch.Tensor)
+    assert loaded_img[0].shape == (1, img_size[1], img_size[0], 3)  # (batch, height, width, channels)
+
+    # Test saving an image
+    save_node = PathSaveImageRGB()
+    output_path = str(tmp_path / "output_rgb")
+
+    # Create a simple red test image tensor
+    red_img = torch.zeros(1, img_size[1], img_size[0], 3)
+    red_img[0, :, :, 0] = 1.0  # Red channel set to 1
+
+    # Save the image
+    assert save_node.save_image(red_img, output_path) == (True,)
+
+    # Verify the image was saved with the correct extension
+    assert os.path.exists(output_path + ".png")
+
+    # Test different formats
+    assert save_node.save_image(red_img, str(tmp_path / "jpeg_test"), format="jpg") == (True,)
+    assert os.path.exists(str(tmp_path / "jpeg_test.jpg"))
+
+    # Test with directories that don't exist
+    nested_path = str(tmp_path / "nested" / "images" / "test_rgb")
+    assert save_node.save_image(red_img, nested_path) == (True,)
+    assert os.path.exists(nested_path + ".png")
+
+    # Test error handling for loading
+    with pytest.raises(FileNotFoundError):
+        load_node.load_image_rgb(str(tmp_path / "nonexistent.png"))
+
+
 
 def test_path_normalize():
     node = PathNormalize()