| import torch | |
| import hashlib | |
| from pathlib import Path | |
| from typing import Iterable | |
| from PIL import Image, ImageOps | |
| import numpy as np | |
| import folder_paths | |
| class LoadImageFromPath: | |
| def INPUT_TYPES(s): | |
| return {"required": | |
| {"image": ("STRING", {"default": r"ComfyUI_00001_-assets\ComfyUI_00001_.png [output]"})}, | |
| } | |
| CATEGORY = "image" | |
| RETURN_TYPES = ("IMAGE", "MASK") | |
| FUNCTION = "load_image" | |
| def load_image(self, image): | |
| image_path = LoadImageFromPath._resolve_path(image) | |
| i = Image.open(image_path) | |
| i = ImageOps.exif_transpose(i) | |
| image = i.convert("RGB") | |
| image = np.array(image).astype(np.float32) / 255.0 | |
| image = torch.from_numpy(image)[None,] | |
| if 'A' in i.getbands(): | |
| mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 | |
| mask = 1. - torch.from_numpy(mask) | |
| else: | |
| mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") | |
| return (image, mask) | |
| def _resolve_path(image) -> Path: | |
| image_path = Path(folder_paths.get_annotated_filepath(image)) | |
| return image_path | |
| def IS_CHANGED(s, image): | |
| image_path = LoadImageFromPath._resolve_path(image) | |
| m = hashlib.sha256() | |
| with open(image_path, 'rb') as f: | |
| m.update(f.read()) | |
| return m.digest().hex() | |
| def VALIDATE_INPUTS(s, image): | |
| # If image is an output of another node, it will be None during validation | |
| if image is None: | |
| return True | |
| image_path = LoadImageFromPath._resolve_path(image) | |
| if not image_path.exists(): | |
| return "Invalid image path: {}".format(image_path) | |
| return True | |
| class PILToImage: | |
| def INPUT_TYPES(s): | |
| return {'required': | |
| {'images': ('PIL_IMAGE', )}, | |
| } | |
| RETURN_TYPES = ('IMAGE',) | |
| FUNCTION = 'pil_images_to_images' | |
| CATEGORY = 'image/PIL' | |
| def pil_images_to_images(images: Iterable[Image.Image]) -> torch.Tensor: | |
| pil_images = images | |
| images = [] | |
| for pil_image in pil_images: | |
| i = pil_image | |
| i = ImageOps.exif_transpose(i) | |
| if i.mode == 'I': | |
| i = i.point(lambda i: i * (1 / 255)) | |
| image = i.convert("RGB") | |
| image = np.array(image).astype(np.float32) / 255.0 | |
| image = torch.from_numpy(image)[None,] | |
| images.append(image) | |
| if len(images) > 1: | |
| images = torch.cat(images, dim=0) | |
| else: | |
| images = images[0] | |
| return (images,) | |
| class PILToMask: | |
| def INPUT_TYPES(s): | |
| return {'required': | |
| {'images': ('PIL_IMAGE', )}, | |
| } | |
| RETURN_TYPES = ('IMAGE',) | |
| FUNCTION = 'pil_images_to_masks' | |
| CATEGORY = 'image/PIL' | |
| def pil_images_to_masks(images: Iterable[Image.Image]) -> torch.Tensor: | |
| pil_images = images | |
| masks = [] | |
| for pil_image in pil_images: | |
| i = pil_image | |
| i = ImageOps.exif_transpose(i) | |
| if i.mode == 'I': | |
| i = i.point(lambda i: i * (1 / 255)) | |
| if 'A' in i.getbands(): | |
| mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 | |
| mask = 1. - torch.from_numpy(mask) | |
| else: | |
| mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") | |
| masks.append(mask) | |
| if len(masks) > 1: | |
| masks = torch.cat(masks, dim=0) | |
| else: | |
| masks = masks[0] | |
| return (masks,) | |
| class ImageToPIL: | |
| def INPUT_TYPES(s): | |
| return {'required': | |
| {'images': ('IMAGE', )}, | |
| } | |
| RETURN_TYPES = ('PIL_IMAGE',) | |
| FUNCTION = 'images_to_pil_images' | |
| CATEGORY = 'image/PIL' | |
| def images_to_pil_images(self, images: torch.Tensor) -> list[Image.Image]: | |
| pil_images = [] | |
| for image in images: | |
| i = 255. * image.cpu().numpy() | |
| img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) | |
| pil_images.append(img) | |
| return (pil_images,) |