fix bugs, ImageAutoCrop add cropped_mask output.

This commit is contained in:
chflame163
2024-02-25 18:26:58 +08:00
committed by GitHub
parent 04001f797f
commit 92b25d9efe
3 changed files with 317 additions and 147 deletions
+157 -146
View File
@@ -1,147 +1,158 @@
from .imagefunc import *
from .segment_anything_func import *
NODE_NAME = 'ImageAutoCrop'
class ImageAutoCrop:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(self):
matting_method_list = ['RMBG 1.4', 'SegmentAnything']
detect_mode = ['min_bounding_rect', 'max_inscribed_rect']
ratio_list = ['1:1', '3:2', '4:3', '16:9', '2:3', '3:4', '9:16', 'custom']
return {
"required": {
"image": ("IMAGE", ), #
"background_color": ("STRING", {"default": "#FFFFFF"}), # 背景颜色
"aspect_ratio": (ratio_list,),
"proportional_width": ("INT", {"default": 2, "min": 1, "max": 999, "step": 1}),
"proportional_height": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}),
"scale_to_longest_side": ("BOOLEAN", {"default": True}), # 是否按长边缩放
"longest_side": ("INT", {"default": 1024, "min": 4, "max": 999999, "step": 1}),
"detect": (detect_mode,),
"border_reserve": ("INT", {"default": 100, "min": -9999, "max": 9999, "step": 1}),
"ultra_detail_range": ("INT", {"default": 0, "min": 0, "max": 256, "step": 1}),
"matting_method": (matting_method_list,),
"sam_model": (list_sam_model(),),
"grounding_dino_model": (list_groundingdino_model(),),
"sam_threshold": ("FLOAT", {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01}),
"sam_prompt": ("STRING", {"default": "subject"}),
},
"optional": {
}
}
RETURN_TYPES = ("IMAGE", "IMAGE",)
RETURN_NAMES = ("cropped_image", "box_preview")
FUNCTION = 'image_auto_crop'
CATEGORY = '😺dzNodes/LayerUtility'
OUTPUT_NODE = True
def image_auto_crop(self, image, detect, border_reserve, aspect_ratio, proportional_width, proportional_height,
background_color, ultra_detail_range, scale_to_longest_side, longest_side,
matting_method, sam_model, grounding_dino_model, sam_threshold, sam_prompt
):
ret_images = []
ret_box_previews = []
input_images = []
input_masks = []
crop_boxs = []
for l in image:
input_images.append(torch.unsqueeze(l, 0))
m = tensor2pil(l)
if m.mode == 'RGBA':
input_masks.append(m.split()[-1])
if len(input_masks) > 0 and len(input_masks) != len(input_images):
input_masks = []
log(f"Warning, {NODE_NAME} unable align alpha to image, drop it.", message_type='warning')
if aspect_ratio == 'custom':
ratio = proportional_width / proportional_height
else:
s = aspect_ratio.split(":")
ratio = int(s[0]) / int(s[1])
side_limit = longest_side if scale_to_longest_side else 0
for i in range(len(input_images)):
_image = tensor2pil(input_images[i]).convert('RGB')
if len(input_masks) > 0:
_mask = input_masks[i]
else:
if matting_method == 'SegmentAnything':
sam_model = load_sam_model(sam_model)
dino_model = load_groundingdino_model(grounding_dino_model)
item = _image.convert('RGBA')
boxes = groundingdino_predict(dino_model, item, sam_prompt, sam_threshold)
(_, _mask) = sam_segment(sam_model, item, boxes)
_mask = mask2image(_mask[0])
else:
_mask = RMBG(_image)
if ultra_detail_range:
_mask = tensor2pil(mask_edge_detail(input_images[i], pil2tensor(_mask), ultra_detail_range, 0.01, 0.99))
bluredmask = gaussian_blur(_mask, 20).convert('L')
x = 0
y = 0
width = 0
height = 0
x_offset = 0
y_offset = 0
if detect == "min_bounding_rect":
(x, y, width, height) = min_bounding_rect(bluredmask)
if detect == "max_inscribed_rect":
(x, y, width, height) = max_inscribed_rect(bluredmask)
canvas_width, canvas_height = _image.size
x1 = x - border_reserve
y1 = y - border_reserve
x2 = x + width + border_reserve
y2 = y + height + border_reserve
if x1 < 0:
canvas_width -= x1
x_offset = -x1
if y1 < 0:
canvas_height -= y1
y_offset = -y1
if x2 > _image.width:
canvas_width += x2 - _image.width
if y2 > _image.height:
canvas_height += y2 - _image.height
crop_box = (x1 + x_offset, y1 + y_offset, width + border_reserve*2, height + border_reserve*2)
crop_boxs.append(crop_box)
if len(crop_boxs) > 0: # 批量图强制使用同一尺寸
crop_box = crop_boxs[0]
target_width, target_height = calculate_side_by_ratio(crop_box[2], crop_box[3], ratio,
longest_side=side_limit)
_canvas = Image.new('RGB', size=(canvas_width, canvas_height), color=background_color)
if ultra_detail_range:
_image = pixel_spread(_image, _mask)
_canvas.paste(_image, box=(x_offset, y_offset), mask=_mask.convert('L'))
preview_image = Image.new('RGB', size=(canvas_width, canvas_height), color='gray')
preview_image.paste(_mask, box=(x_offset, y_offset))
preview_image = draw_rect(preview_image,
crop_box[0], crop_box[1], crop_box[2], crop_box[3],
line_color="#F00000", line_width=(canvas_width + canvas_height)//200)
ret_image = _canvas.crop((crop_box[0], crop_box[1], crop_box[0]+crop_box[2], crop_box[1]+crop_box[3]))
ret_image = fit_resize_image(ret_image, target_width, target_height,
fit='letterbox', resize_sampler=Image.LANCZOS,
background_color=background_color)
ret_images.append(pil2tensor(ret_image))
ret_box_previews.append(pil2tensor(preview_image))
log(f"{NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0), torch.cat(ret_box_previews, dim=0),)
NODE_CLASS_MAPPINGS = {
"LayerUtility: ImageAutoCrop": ImageAutoCrop
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerUtility: ImageAutoCrop": "LayerUtility: ImageAutoCrop"
from .imagefunc import *
from .segment_anything_func import *
NODE_NAME = 'ImageAutoCrop'
class ImageAutoCrop:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(self):
matting_method_list = ['RMBG 1.4', 'SegmentAnything']
detect_mode = ['min_bounding_rect', 'max_inscribed_rect']
ratio_list = ['1:1', '3:2', '4:3', '16:9', '2:3', '3:4', '9:16', 'custom']
return {
"required": {
"image": ("IMAGE", ), #
"background_color": ("STRING", {"default": "#FFFFFF"}), # 背景颜色
"aspect_ratio": (ratio_list,),
"proportional_width": ("INT", {"default": 2, "min": 1, "max": 999, "step": 1}),
"proportional_height": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}),
"scale_to_longest_side": ("BOOLEAN", {"default": True}), # 是否按长边缩放
"longest_side": ("INT", {"default": 1024, "min": 4, "max": 999999, "step": 1}),
"detect": (detect_mode,),
"border_reserve": ("INT", {"default": 100, "min": -9999, "max": 9999, "step": 1}),
"ultra_detail_range": ("INT", {"default": 0, "min": 0, "max": 256, "step": 1}),
"matting_method": (matting_method_list,),
"sam_model": (list_sam_model(),),
"grounding_dino_model": (list_groundingdino_model(),),
"sam_threshold": ("FLOAT", {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01}),
"sam_prompt": ("STRING", {"default": "subject"}),
},
"optional": {
}
}
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK",)
RETURN_NAMES = ("cropped_image", "box_preview", "cropped_mask",)
FUNCTION = 'image_auto_crop'
CATEGORY = '😺dzNodes/LayerUtility'
OUTPUT_NODE = True
def image_auto_crop(self, image, detect, border_reserve, aspect_ratio, proportional_width, proportional_height,
background_color, ultra_detail_range, scale_to_longest_side, longest_side,
matting_method, sam_model, grounding_dino_model, sam_threshold, sam_prompt
):
ret_images = []
ret_box_previews = []
ret_masks = []
input_images = []
input_masks = []
crop_boxs = []
for l in image:
input_images.append(torch.unsqueeze(l, 0))
m = tensor2pil(l)
if m.mode == 'RGBA':
input_masks.append(m.split()[-1])
if len(input_masks) > 0 and len(input_masks) != len(input_images):
input_masks = []
log(f"Warning, {NODE_NAME} unable align alpha to image, drop it.", message_type='warning')
if aspect_ratio == 'custom':
ratio = proportional_width / proportional_height
else:
s = aspect_ratio.split(":")
ratio = int(s[0]) / int(s[1])
side_limit = longest_side if scale_to_longest_side else 0
for i in range(len(input_images)):
_image = tensor2pil(input_images[i]).convert('RGB')
if len(input_masks) > 0:
_mask = input_masks[i]
else:
if matting_method == 'SegmentAnything':
sam_model = load_sam_model(sam_model)
dino_model = load_groundingdino_model(grounding_dino_model)
item = _image.convert('RGBA')
boxes = groundingdino_predict(dino_model, item, sam_prompt, sam_threshold)
(_, _mask) = sam_segment(sam_model, item, boxes)
_mask = mask2image(_mask[0])
else:
_mask = RMBG(_image)
if ultra_detail_range:
_mask = tensor2pil(mask_edge_detail(input_images[i], pil2tensor(_mask), ultra_detail_range, 0.01, 0.99))
bluredmask = gaussian_blur(_mask, 20).convert('L')
x = 0
y = 0
width = 0
height = 0
x_offset = 0
y_offset = 0
if detect == "min_bounding_rect":
(x, y, width, height) = min_bounding_rect(bluredmask)
if detect == "max_inscribed_rect":
(x, y, width, height) = max_inscribed_rect(bluredmask)
canvas_width, canvas_height = _image.size
x1 = x - border_reserve
y1 = y - border_reserve
x2 = x + width + border_reserve
y2 = y + height + border_reserve
if x1 < 0:
canvas_width -= x1
x_offset = -x1
if y1 < 0:
canvas_height -= y1
y_offset = -y1
if x2 > _image.width:
canvas_width += x2 - _image.width
if y2 > _image.height:
canvas_height += y2 - _image.height
crop_box = (x1 + x_offset, y1 + y_offset, width + border_reserve*2, height + border_reserve*2)
crop_boxs.append(crop_box)
if len(crop_boxs) > 0: # 批量图强制使用同一尺寸
crop_box = crop_boxs[0]
target_width, target_height = calculate_side_by_ratio(crop_box[2], crop_box[3], ratio,
longest_side=side_limit)
_canvas = Image.new('RGB', size=(canvas_width, canvas_height), color=background_color)
_mask_canvas = Image.new('L', size=(canvas_width, canvas_height), color='black')
if ultra_detail_range:
_image = pixel_spread(_image, _mask)
_canvas.paste(_image, box=(x_offset, y_offset), mask=_mask.convert('L'))
_mask_canvas.paste(_mask, box=(x_offset, y_offset))
preview_image = Image.new('RGB', size=(canvas_width, canvas_height), color='gray')
preview_image.paste(_mask, box=(x_offset, y_offset))
preview_image = draw_rect(preview_image,
crop_box[0], crop_box[1], crop_box[2], crop_box[3],
line_color="#F00000", line_width=(canvas_width + canvas_height)//200)
ret_image = _canvas.crop((crop_box[0], crop_box[1], crop_box[0]+crop_box[2], crop_box[1]+crop_box[3]))
ret_image = fit_resize_image(ret_image, target_width, target_height,
fit='letterbox', resize_sampler=Image.LANCZOS,
background_color=background_color)
ret_mask = _mask_canvas.crop((crop_box[0], crop_box[1], crop_box[0]+crop_box[2], crop_box[1]+crop_box[3]))
ret_mask = fit_resize_image(ret_mask, target_width, target_height,
fit='letterbox', resize_sampler=Image.LANCZOS,
background_color="#000000")
ret_images.append(pil2tensor(ret_image))
ret_box_previews.append(pil2tensor(preview_image))
ret_masks.append(image2mask(ret_mask))
log(f"{NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0),
torch.cat(ret_box_previews, dim=0),
torch.cat(ret_masks, dim=0),
)
NODE_CLASS_MAPPINGS = {
"LayerUtility: ImageAutoCrop": ImageAutoCrop
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerUtility: ImageAutoCrop": "LayerUtility: ImageAutoCrop"
}
+158
View File
@@ -0,0 +1,158 @@
import torch
from .imagefunc import *
NODE_NAME = 'ImageScaleByAspectRatio'
class ImageScaleByAspectRatio:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(self):
ratio_list = ['original', 'custom', '1:1', '3:2', '4:3', '16:9', '2:3', '3:4', '9:16']
fit_mode = ['letterbox', 'crop', 'fill']
method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest']
multiple_list = ['8', '16', 'None']
return {
"required": {
"aspect_ratio": (ratio_list,),
"proportional_width": ("INT", {"default": 2, "min": 1, "max": 999, "step": 1}),
"proportional_height": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}),
"fit": (fit_mode,),
"method": (method_mode,),
"round_to_multiple": (multiple_list,),
"scale_to_longest_side": ("BOOLEAN", {"default": False}), # 是否按长边缩放
"longest_side": ("INT", {"default": 1024, "min": 4, "max": 999999, "step": 1}),
},
"optional": {
"image": ("IMAGE",), #
"mask": ("MASK",), #
}
}
RETURN_TYPES = ("IMAGE", "MASK", "BOX",)
RETURN_NAMES = ("image", "mask", "original_size")
FUNCTION = 'image_scale_by_aspect_ratio'
CATEGORY = '😺dzNodes/LayerUtility'
OUTPUT_NODE = True
def image_scale_by_aspect_ratio(self, aspect_ratio, proportional_width, proportional_height,
fit, method, round_to_multiple, scale_to_longest_side, longest_side,
image=None, mask = None,
):
orig_images = []
orig_masks = []
orig_width = 0
orig_height = 0
target_width = 0
target_height = 0
ratio = 1.0
ret_images = []
ret_masks = []
if image is not None:
for i in image:
i = torch.unsqueeze(i, 0)
orig_images.append(i)
orig_width, orig_height = tensor2pil(orig_images[0]).size
if mask is not None:
if mask.dim() == 2:
mask = torch.unsqueeze(mask, 0)
for m in mask:
m = torch.unsqueeze(m, 0)
orig_masks.append(m)
_width, _height = tensor2pil(orig_masks[0]).size
if (orig_width > 0 and orig_width != _width) or (orig_height > 0 and orig_height != _height):
log(f"Error: {NODE_NAME} skipped, because the mask is does'nt match image.", message_type='error')
return (None, None,)
elif orig_width + orig_height == 0:
orig_width = _width
orig_height = _height
if orig_width + orig_height == 0:
log(f"Error: {NODE_NAME} skipped, because the image or mask at least one must be input.", message_type='error')
return (None, None,)
if aspect_ratio == 'original':
ratio = orig_width / orig_height
elif aspect_ratio == 'custom':
ratio = proportional_width / proportional_height
else:
s = aspect_ratio.split(":")
ratio = int(s[0]) / int(s[1])
# calculate target width and height
if orig_width > orig_height:
if scale_to_longest_side:
target_width = longest_side
else:
target_width = orig_width
target_height = int(target_width / ratio)
else:
if scale_to_longest_side:
target_height = longest_side
else:
target_height = orig_height
target_width = int(target_height * ratio)
if ratio < 1:
if scale_to_longest_side:
_r = longest_side / target_height
target_height = longest_side
else:
_r = orig_height / target_height
target_height = orig_height
target_width = int(target_width * _r)
if round_to_multiple != 'None':
multiple = int(round_to_multiple)
target_width = num_round_to_multiple(target_width, multiple)
target_height = num_round_to_multiple(target_height, multiple)
_mask = Image.new('L', size=(target_width, target_height), color='black')
_image = Image.new('RGB', size=(target_width, target_height), color='black')
resize_sampler = Image.LANCZOS
if method == "bicubic":
resize_sampler = Image.BICUBIC
elif method == "hamming":
resize_sampler = Image.HAMMING
elif method == "bilinear":
resize_sampler = Image.BILINEAR
elif method == "box":
resize_sampler = Image.BOX
elif method == "nearest":
resize_sampler = Image.NEAREST
if len(orig_images) > 0:
for i in orig_images:
_image = tensor2pil(i).convert('RGB')
_image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler)
ret_images.append(pil2tensor(_image))
if len(orig_masks) > 0:
for m in orig_masks:
_mask = tensor2pil(m).convert('L')
_mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler).convert('L')
ret_masks.append(image2mask(_mask))
if len(ret_images) > 0 and len(ret_masks) >0:
log(f"{NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),)
elif len(ret_images) > 0 and len(ret_masks) == 0:
log(f"{NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0), None,)
elif len(ret_images) == 0 and len(ret_masks) > 0:
log(f"{NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish')
return (None, torch.cat(ret_masks, dim=0),)
else:
log(f"Error: {NODE_NAME} skipped, because the available image or mask is not found.", message_type='error')
return (None, None,)
NODE_CLASS_MAPPINGS = {
"LayerUtility: ImageScaleByAspectRatio": ImageScaleByAspectRatio
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerUtility: ImageScaleByAspectRatio": "LayerUtility: ImageScaleByAspectRatio"
}
+2 -1
View File
@@ -41,7 +41,8 @@ def log(message:str, message_type:str='info'):
try:
from cv2.ximgproc import guidedFilter
except ImportError:
except ImportError as e:
print(e)
log(f'Dependency package error, unable import "cv2.ximgproc".'
f'\nPlease REINSTALL package "opencv-contrib-python".'
f'\nFor detail refer to \033[4mhttps://github.com/chflame163/ComfyUI_LayerStyle/issues/5\033[0m',