diff --git a/py/image_auto_crop.py b/py/image_auto_crop.py index 727b9ec..e4be434 100644 --- a/py/image_auto_crop.py +++ b/py/image_auto_crop.py @@ -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" } \ No newline at end of file diff --git a/py/image_scale_by_aspect_ratio.py b/py/image_scale_by_aspect_ratio.py new file mode 100644 index 0000000..b013d41 --- /dev/null +++ b/py/image_scale_by_aspect_ratio.py @@ -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" +} \ No newline at end of file diff --git a/py/imagefunc.py b/py/imagefunc.py index 1bd9ec5..d509a8c 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -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',