Spaces:
Runtime error
Runtime error
File size: 2,226 Bytes
f2dbf59 |
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 |
import os
import cv2
import torch
import requests
import itertools
import folder_paths
import psutil
import numpy as np
from comfy.utils import common_upscale
from io import BytesIO
from PIL import Image, ImageSequence, ImageOps
def pil2tensor(img):
output_images = []
output_masks = []
for i in ImageSequence.Iterator(img):
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,]
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")
output_images.append(image)
output_masks.append(mask.unsqueeze(0))
if len(output_images) > 1:
output_image = torch.cat(output_images, dim=0)
output_mask = torch.cat(output_masks, dim=0)
else:
output_image = output_images[0]
output_mask = output_masks[0]
return (output_image, output_mask)
def load_image(image_source):
if image_source.startswith('http'):
print(image_source)
response = requests.get(image_source)
img = Image.open(BytesIO(response.content))
file_name = image_source.split('/')[-1]
else:
img = Image.open(image_source)
file_name = os.path.basename(image_source)
return img, file_name
class LoadImageNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING", {"multiline": True, "dynamicPrompts": False})
}
}
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "load_image"
CATEGORY = "tbox/Image"
def load_image(self, path):
filepaht = path.split('\n')[0]
img, name = load_image(filepaht)
img_out, mask_out = pil2tensor(img)
return (img_out, mask_out)
if __name__ == "__main__":
img, name = load_image("https://creativestorage.blob.core.chinacloudapi.cn/test/bird.png")
img_out, mask_out = pil2tensor(img)
|