File size: 2,511 Bytes
a737419 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import os
import json
import hashlib
import numpy as np
import torch
from PIL import Image, ImageOps
import folder_paths
class TZ_LoadPNGWithMetadata:
@classmethod
def INPUT_TYPES(cls):
input_dir = folder_paths.get_input_directory()
files = []
if os.path.isdir(input_dir):
for f in os.listdir(input_dir):
if f.lower().endswith(".png"):
files.append(f)
files.sort()
if not files:
files = ["example.png"]
return {
"required": {
"image": (files, {"image_upload": True}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING")
RETURN_NAMES = ("image", "mask", "image_path", "raw_metadata")
FUNCTION = "load_image"
CATEGORY = "Herve/metadata"
@classmethod
def IS_CHANGED(cls, image):
image_path = folder_paths.get_annotated_filepath(image)
m = hashlib.sha256()
m.update(image.encode("utf-8"))
if os.path.exists(image_path):
stat = os.stat(image_path)
m.update(str(stat.st_mtime_ns).encode("utf-8"))
m.update(str(stat.st_size).encode("utf-8"))
return m.hexdigest()
@classmethod
def VALIDATE_INPUTS(cls, image):
if not image.lower().endswith(".png"):
return "Este nodo solo acepta archivos PNG"
image_path = folder_paths.get_annotated_filepath(image)
if not os.path.isfile(image_path):
return f"No se encontró el archivo: {image}"
return True
def load_image(self, image):
image_path = folder_paths.get_annotated_filepath(image)
img = Image.open(image_path)
img = ImageOps.exif_transpose(img)
info = getattr(img, "info", {}) or {}
raw_metadata = json.dumps(info, ensure_ascii=False, indent=2)
image_rgb = img.convert("RGB")
image_np = np.array(image_rgb).astype(np.float32) / 255.0
image_tensor = torch.from_numpy(image_np)[None,]
if "A" in img.getbands():
mask_np = np.array(img.getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask_np)
else:
mask = torch.zeros((64, 64), dtype=torch.float32)
return (image_tensor, mask, image_path, raw_metadata)
NODE_CLASS_MAPPINGS = {
"TZ_LoadPNGWithMetadata": TZ_LoadPNGWithMetadata
}
NODE_DISPLAY_NAME_MAPPINGS = {
"TZ_LoadPNGWithMetadata": "TZ Load PNG With Metadata"
} |