65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
# pk_load_image_paste v7 - node implementation using STRING input to avoid 'Value not in list' issues
|
|
import os
|
|
import folder_paths
|
|
|
|
class PK_LoadImage_Paste:
|
|
"""Load an image (path relative to input/ or absolute) and output (IMAGE, MASK)."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("STRING", {
|
|
"default": "",
|
|
"placeholder": "e.g. pasted/xxx.png or input/xxx.png",
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "load"
|
|
CATEGORY = "image"
|
|
|
|
def load(self, image: str):
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
import torch
|
|
|
|
if not image or not isinstance(image, str):
|
|
raise ValueError("image path is empty")
|
|
|
|
base = folder_paths.get_input_directory()
|
|
# allow absolute or relative-to-input paths
|
|
if os.path.isabs(image):
|
|
path = image
|
|
else:
|
|
# strip leading separators or 'input/' prefix for safety
|
|
sanitized = image.strip().lstrip("/\\")
|
|
if sanitized.startswith("input/") or sanitized.startswith("input\\"):
|
|
sanitized = sanitized.split("/", 1)[1] if "/" in sanitized else sanitized.split("\\", 1)[1]
|
|
path = os.path.join(base, sanitized)
|
|
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Image not found: {path}")
|
|
|
|
img = Image.open(path)
|
|
img = ImageOps.exif_transpose(img)
|
|
|
|
if img.mode == "I;16":
|
|
img = img.point(lambda i: i * (1 / 257)).convert("L")
|
|
|
|
if img.mode == "RGBA":
|
|
a = img.getchannel("A")
|
|
mask_np = 1.0 - (np.asarray(a, dtype=np.float32) / 255.0)
|
|
img = img.convert("RGB")
|
|
else:
|
|
mask_np = np.zeros((img.height, img.width), dtype=np.float32)
|
|
|
|
arr = np.asarray(img, dtype=np.float32) / 255.0
|
|
if arr.ndim == 2:
|
|
arr = np.expand_dims(arr, 2)
|
|
|
|
image_t = torch.from_numpy(arr)[None, ...]
|
|
mask_t = torch.from_numpy(mask_np)[None, None, ...]
|
|
return (image_t, mask_t)
|