109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
# Helper functions that specifically involve images.
|
|
|
|
import io
|
|
import base64
|
|
import time
|
|
import pathlib
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
import numpy as np
|
|
import torch
|
|
import requests
|
|
|
|
import node_helpers
|
|
import folder_paths
|
|
|
|
|
|
def blank_image():
|
|
"""Create a blank 1024x1024 RGB image as a torch tensor."""
|
|
img = Image.new('RGB', (1024, 1024))
|
|
img = ImageOps.exif_transpose(img)
|
|
img = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
|
return torch.from_numpy(img)[None, :]
|
|
|
|
|
|
def url_to_torch_image(url):
|
|
"""Load an image from a URL and return as a torch tensor."""
|
|
response = requests.get(url, stream=True)
|
|
img = Image.open(io.BytesIO(response.content))
|
|
img = ImageOps.exif_transpose(img)
|
|
img = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
|
return torch.from_numpy(img)[None, :]
|
|
|
|
|
|
def tensor_to_base64(tensor):
|
|
"""Convert a torch tensor image batch to a list of base64-encoded PNGs."""
|
|
if tensor is None or not isinstance(tensor, torch.Tensor):
|
|
return []
|
|
base64_images = []
|
|
for image in tensor:
|
|
arr = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
|
|
if img.mode != 'RGB':
|
|
img = img.convert('RGB')
|
|
buffered = io.BytesIO()
|
|
img.save(buffered, format="PNG")
|
|
base64_images.append(base64.b64encode(buffered.getvalue()).decode('utf-8'))
|
|
return base64_images
|
|
|
|
|
|
def _load_image(img) -> tuple:
|
|
"""Internal helper to process a PIL image (from path or URL) into torch tensors and mask."""
|
|
output_images, output_masks = [], []
|
|
w, h = None, None
|
|
for i in ImageSequence.Iterator(img):
|
|
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
|
if i.mode == "I":
|
|
i = i.point(lambda x: x * (1 / 255))
|
|
image = i.convert("RGB")
|
|
if not output_images:
|
|
w, h = image.size
|
|
if image.size != (w, h):
|
|
continue
|
|
image_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None, :]
|
|
if "A" in i.getbands():
|
|
mask = (1.0 - torch.from_numpy(np.array(i.getchannel("A")).astype(np.float32) / 255.0)).unsqueeze(0)
|
|
else:
|
|
mask = torch.zeros((1, 64, 64), dtype=torch.float32)
|
|
output_images.append(image_tensor)
|
|
output_masks.append(mask)
|
|
output_image = torch.cat(output_images, dim=0) if len(output_images) > 1 and getattr(img, 'format', None) != "MPO" else output_images[0]
|
|
output_mask = torch.cat(output_masks, dim=0) if len(output_masks) > 1 and getattr(img, 'format', None) != "MPO" else output_masks[0]
|
|
return output_image, output_mask, w, h, f"{getattr(img, 'info', {})}"
|
|
|
|
|
|
def load_image_from_path(image_path) -> tuple:
|
|
"""Load an image (and mask if present) from a file path as torch tensors."""
|
|
img = node_helpers.pillow(Image.open, image_path)
|
|
return _load_image(img)
|
|
|
|
|
|
def load_image_from_url(url) -> tuple:
|
|
"""Load an image (and mask if present) from a URL as torch tensors."""
|
|
response = requests.get(url, stream=True)
|
|
img = node_helpers.pillow(Image.open, io.BytesIO(response.content))
|
|
return _load_image(img)
|
|
|
|
|
|
def tensor_to_temp_image(tensor, filename=None):
|
|
"""Save a torch tensor image batch to temporary PNG files. Returns list of file paths."""
|
|
if tensor is None or not isinstance(tensor, torch.Tensor):
|
|
return []
|
|
output_dir = folder_paths.get_temp_directory()
|
|
filenames = []
|
|
for idx, image in enumerate(tensor):
|
|
arr = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
|
|
if img.mode != 'RGB':
|
|
img = img.convert('RGB')
|
|
if filename is None:
|
|
fname = f"temp_image_{idx}_{int(time.time())}.png"
|
|
else:
|
|
stem = pathlib.Path(filename).stem
|
|
fname = f"{stem}_{idx}.png"
|
|
file_path = pathlib.Path(output_dir) / fname
|
|
img.save(file_path, format="PNG")
|
|
filenames.append(str(file_path))
|
|
print(f"Saved {len(filenames)} images to {output_dir}")
|
|
print(filenames)
|
|
return filenames
|