148 lines
4.2 KiB
Python
148 lines
4.2 KiB
Python
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:
|
|
@classmethod
|
|
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
|
|
|
|
@classmethod
|
|
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()
|
|
|
|
@classmethod
|
|
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:
|
|
@classmethod
|
|
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:
|
|
@classmethod
|
|
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:
|
|
@classmethod
|
|
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,) |