312 lines
7.8 KiB
Python
312 lines
7.8 KiB
Python
from PIL import Image
|
|
from enum import Enum
|
|
import os
|
|
import hashlib
|
|
import folder_paths
|
|
import torch
|
|
import numpy as np
|
|
from nodes import MAX_RESOLUTION
|
|
from .log import log_node_info
|
|
|
|
class ResizeMode(Enum):
|
|
RESIZE = "Just Resize"
|
|
INNER_FIT = "Crop and Resize"
|
|
OUTER_FIT = "Resize and Fill"
|
|
def int_value(self):
|
|
if self == ResizeMode.RESIZE:
|
|
return 0
|
|
elif self == ResizeMode.INNER_FIT:
|
|
return 1
|
|
elif self == ResizeMode.OUTER_FIT:
|
|
return 2
|
|
assert False, "NOTREACHED"
|
|
|
|
RESIZE_MODES = [ResizeMode.RESIZE.value, ResizeMode.INNER_FIT.value, ResizeMode.OUTER_FIT.value]
|
|
|
|
def get_new_bounds(width, height, left, right, top, bottom):
|
|
"""Returns the new bounds for an image with inset crop data."""
|
|
left = 0 + left
|
|
right = width - right
|
|
top = 0 + top
|
|
bottom = height - bottom
|
|
return (left, right, top, bottom)
|
|
|
|
# Tensor to PIL
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
# 图像裁切
|
|
class imageInsetCrop:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls): # pylint: disable = invalid-name, missing-function-docstring
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"measurement": (['Pixels', 'Percentage'],),
|
|
"left": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": MAX_RESOLUTION,
|
|
"step": 8
|
|
}),
|
|
"right": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": MAX_RESOLUTION,
|
|
"step": 8
|
|
}),
|
|
"top": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": MAX_RESOLUTION,
|
|
"step": 8
|
|
}),
|
|
"bottom": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": MAX_RESOLUTION,
|
|
"step": 8
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "crop"
|
|
|
|
CATEGORY = "EasyUse/Image"
|
|
|
|
# pylint: disable = too-many-arguments
|
|
def crop(self, measurement, left, right, top, bottom, image=None):
|
|
"""Does the crop."""
|
|
|
|
_, height, width, _ = image.shape
|
|
|
|
if measurement == 'Percentage':
|
|
left = int(width - (width * (100 - left) / 100))
|
|
right = int(width - (width * (100 - right) / 100))
|
|
top = int(height - (height * (100 - top) / 100))
|
|
bottom = int(height - (height * (100 - bottom) / 100))
|
|
|
|
# Snap to 8 pixels
|
|
left = left // 8 * 8
|
|
right = right // 8 * 8
|
|
top = top // 8 * 8
|
|
bottom = bottom // 8 * 8
|
|
|
|
if left == 0 and right == 0 and bottom == 0 and top == 0:
|
|
return (image,)
|
|
|
|
inset_left, inset_right, inset_top, inset_bottom = get_new_bounds(width, height, left, right,
|
|
top, bottom)
|
|
if inset_top > inset_bottom:
|
|
raise ValueError(
|
|
f"Invalid cropping dimensions top ({inset_top}) exceeds bottom ({inset_bottom})")
|
|
if inset_left > inset_right:
|
|
raise ValueError(
|
|
f"Invalid cropping dimensions left ({inset_left}) exceeds right ({inset_right})")
|
|
|
|
log_node_info("Image Inset Crop", f'Cropping image {width}x{height} width inset by {inset_left},{inset_right}, ' +
|
|
f'and height inset by {inset_top}, {inset_bottom}')
|
|
image = image[:, inset_top:inset_bottom, inset_left:inset_right, :]
|
|
|
|
return (image,)
|
|
|
|
# 图像尺寸
|
|
class imageSize:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "INT")
|
|
RETURN_NAMES = ("width_int", "height_int")
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "image_width_height"
|
|
|
|
CATEGORY = "EasyUse/Image"
|
|
|
|
def image_width_height(self, image):
|
|
_, raw_H, raw_W, _ = image.shape
|
|
|
|
width = raw_W
|
|
height = raw_H
|
|
|
|
if width is not None and height is not None:
|
|
result = (width, height)
|
|
else:
|
|
result = (0, 0)
|
|
return {"ui": {"text": "Width: "+str(width)+" , Height: "+str(height)}, "result": result}
|
|
|
|
# 图像尺寸（最长边）
|
|
class imageSizeBySide:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"side": (["Longest", "Shortest"],)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("resolution",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "image_side"
|
|
|
|
CATEGORY = "EasyUse/Image"
|
|
|
|
def image_side(self, image, side):
|
|
_, raw_H, raw_W, _ = image.shape
|
|
|
|
width = raw_W
|
|
height = raw_H
|
|
if width is not None and height is not None:
|
|
if side == "Longest":
|
|
result = (width,) if width > height else (height,)
|
|
elif side == 'Shortest':
|
|
result = (width,) if width < height else (height,)
|
|
else:
|
|
result = (0,)
|
|
return {"ui": {"text": str(result[0])}, "result": result}
|
|
|
|
# 图像尺寸（最长边）
|
|
class imageSizeByLongerSide:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("resolution",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "image_longer_side"
|
|
|
|
CATEGORY = "EasyUse/Image"
|
|
|
|
def image_longer_side(self, image):
|
|
_, raw_H, raw_W, _ = image.shape
|
|
|
|
width = raw_W
|
|
height = raw_H
|
|
if width is not None and height is not None:
|
|
if width > height:
|
|
result = (width,)
|
|
else:
|
|
result = (height,)
|
|
else:
|
|
result = (0,)
|
|
return {"ui": {"text": str(result[0])}, "result": result}
|
|
|
|
# 图像完美像素
|
|
class imagePixelPerfect:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"resize_mode": (RESIZE_MODES, {"default": ResizeMode.RESIZE.value})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("resolution",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "EasyUse/Image"
|
|
|
|
def execute(self, image, resize_mode):
|
|
|
|
_, raw_H, raw_W, _ = image.shape
|
|
|
|
width = raw_W
|
|
height = raw_H
|
|
|
|
k0 = float(height) / float(raw_H)
|
|
k1 = float(width) / float(raw_W)
|
|
|
|
if resize_mode == ResizeMode.OUTER_FIT.value:
|
|
estimation = min(k0, k1) * float(min(raw_H, raw_W))
|
|
else:
|
|
estimation = max(k0, k1) * float(min(raw_H, raw_W))
|
|
|
|
result = int(np.round(estimation))
|
|
text = f"Width:{str(width)}\nHeight:{str(height)}\nPixelPerfect:{str(result)}"
|
|
|
|
return {"ui": {"text": text}, "result": (result,)}
|
|
|
|
# 姿势编辑器
|
|
class poseEditor:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
|
|
if not os.path.isdir(temp_dir):
|
|
os.makedirs(temp_dir)
|
|
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
|
|
return {"required":
|
|
{"image": (sorted(os.listdir(temp_dir)),)},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "output_pose"
|
|
|
|
CATEGORY = "EasyUse/Image"
|
|
|
|
def output_pose(self, image):
|
|
image_path = os.path.join(folder_paths.get_temp_directory(), image)
|
|
# print(f"Create: {image_path}")
|
|
|
|
i = Image.open(image_path)
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
return (image,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, image):
|
|
image_path = os.path.join(
|
|
folder_paths.get_temp_directory(), image)
|
|
# print(f'Change: {image_path}')
|
|
|
|
m = hashlib.sha256()
|
|
with open(image_path, 'rb') as f:
|
|
m.update(f.read())
|
|
return m.digest().hex()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"easy imageInsetCrop": imageInsetCrop,
|
|
"easy imageSize": imageSize,
|
|
"easy imageSizeBySide": imageSizeBySide,
|
|
"easy imageSizeByLongerSide": imageSizeByLongerSide,
|
|
"easy imagePixelPerfect": imagePixelPerfect,
|
|
"easy poseEditor": poseEditor
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"easy imageInsetCrop": "ImageInsetCrop",
|
|
"easy imageSize": "ImageSize",
|
|
"easy imageSizeBySide": "ImageSize (Side)",
|
|
"easy imageSizeByLongerSide": "ImageSize (LongerSide)",
|
|
"easy imagePixelPerfect": "ImagePixelPerfect",
|
|
"easy poseEditor": "PoseEditor"
|
|
} |