251 lines
9.7 KiB
Python
251 lines
9.7 KiB
Python
import torch
|
|
import random
|
|
import cv2
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
from PIL import Image, ImageDraw, ImageFilter
|
|
import copy
|
|
from .imagefunc import log, tensor2pil, pil2tensor, image2mask
|
|
from .imagefunc import fit_resize_image, extract_numbers, gaussian_blur, mask_area, draw_rounded_rectangle
|
|
|
|
class LS_CollageGenerator:
|
|
"""
|
|
随机分割生成指定数量的不规则小矩形。
|
|
"""
|
|
def __init__(self, width, height, num, border_width, r, uniformity, seed ):
|
|
self.width = width
|
|
self.height = height
|
|
self.num = num
|
|
self.border_width = int((self.width + self.height) * border_width / 200)
|
|
self.r = r
|
|
self.seed = seed
|
|
self.split_num = int(1e18)
|
|
self.uniformity = uniformity
|
|
self.rectangles = self.adjust_bboxes_with_gaps(self.split_rec())
|
|
|
|
def split_rec(self):
|
|
random.seed(self.seed)
|
|
if self.num <= 0 or self.width <= 0 or self.height <= 0:
|
|
raise ValueError("Value mast be positive integer")
|
|
|
|
current_rectangles = [(0, 0, self.width, self.height, 0)]
|
|
|
|
while len(current_rectangles) < self.num:
|
|
split_counts = [rect[4] for rect in current_rectangles]
|
|
min_splits = min(split_counts)
|
|
max_splits = max(split_counts)
|
|
probabilities = []
|
|
|
|
for rect in current_rectangles:
|
|
split_count = rect[4]
|
|
normalized_splits = (split_count - min_splits) / (
|
|
max_splits - min_splits if max_splits > min_splits else 1)
|
|
probability = 1 - (normalized_splits * (1 - self.uniformity))
|
|
probabilities.append(probability)
|
|
if sum(probabilities) > 0:
|
|
probabilities = [p / sum(probabilities) for p in probabilities]
|
|
else:
|
|
probabilities = [1.0 / len(probabilities)] * len(probabilities)
|
|
|
|
rect_index = random.choices(range(len(current_rectangles)),
|
|
weights=probabilities, k=1)[0]
|
|
|
|
x, y, w, h, split_count = current_rectangles.pop(rect_index)
|
|
|
|
if w > h or (w == h and random.choice([True, False])):
|
|
split = random.uniform(0.3, 0.7) * w
|
|
rect1 = (x, y, split, h, split_count + 1)
|
|
rect2 = (x + split, y, w - split, h, split_count + 1)
|
|
else:
|
|
split = random.uniform(0.3, 0.7) * h
|
|
rect1 = (x, y, w, split, split_count + 1)
|
|
rect2 = (x, y + split, w, h - split, split_count + 1)
|
|
|
|
current_rectangles.extend([rect1, rect2])
|
|
|
|
rectangles = [(int(x), int(y), int(w), int(h))
|
|
for x, y, w, h, _ in current_rectangles]
|
|
|
|
return rectangles
|
|
|
|
def adjust_bboxes_with_gaps(self, rectangles):
|
|
MIN_SIZE = 1
|
|
adjusted_bboxes = []
|
|
|
|
for x, y, w, h in rectangles:
|
|
new_x = min(x + self.border_width, self.width - MIN_SIZE)
|
|
new_y = min(y + self.border_width, self.height - MIN_SIZE)
|
|
new_w = max(MIN_SIZE, w - 2 * self.border_width)
|
|
new_h = max(MIN_SIZE, h - 2 * self.border_width)
|
|
|
|
if new_x + new_w > self.width:
|
|
new_x = max(0, self.width - new_w)
|
|
if new_y + new_h > self.height:
|
|
new_y = max(0, self.height - new_h)
|
|
|
|
adjusted_bboxes.append((new_x, new_y, new_w, new_h))
|
|
|
|
return adjusted_bboxes
|
|
|
|
def draw_mask(self):
|
|
bboxes = []
|
|
|
|
for bbox in self.rectangles:
|
|
bboxes.append((bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]))
|
|
scale_factor = 2
|
|
|
|
img = Image.new('RGB', (self.width, self.height), color='white')
|
|
img = draw_rounded_rectangle(img, self.r, bboxes, scale_factor, color='black')
|
|
|
|
return img
|
|
|
|
|
|
class LS_Collage:
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Collage'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"canvas_width": ("INT", {"default": 2048, "min": 512, "max": 8192, "step": 16}),
|
|
"canvas_height": ("INT", {"default": 2048, "min": 512, "max": 8192, "step": 16}),
|
|
"border_width": ("FLOAT", {"default": 2, "min": 0, "max": 20, "step": 0.1}),
|
|
"rounded_rect_radius": ("INT", {"default": 8, "min": 0, "max": 100, "step": 1}),
|
|
"uniformity": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.1}), # 分割均匀权重 0均匀分割，1不均匀分割
|
|
"background_color": ("STRING", {"default": "#FFFFFF"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 1e18, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"florence2_model": ("FLORENCE2",),
|
|
"object_prompt": ("STRING", {"default": "face"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK",)
|
|
RETURN_NAMES = ("image", "mask",)
|
|
FUNCTION = "collage"
|
|
CATEGORY = '😺dzNodes/LayerUtility'
|
|
|
|
def collage(self, images, canvas_width, canvas_height, border_width, rounded_rect_radius,
|
|
uniformity, background_color, seed, florence2_model=None, object_prompt="face"):
|
|
|
|
batch_size = images.shape[0]
|
|
|
|
rects = LS_CollageGenerator(width=canvas_width,
|
|
height=canvas_height,
|
|
num=batch_size,
|
|
border_width=border_width,
|
|
r=rounded_rect_radius,
|
|
uniformity=uniformity,
|
|
seed=seed)
|
|
|
|
rects_border_image = rects.draw_mask()
|
|
canvas = Image.new("RGB", (canvas_width, canvas_height), color=background_color)
|
|
color_image = copy.deepcopy(canvas)
|
|
from .object_detector import LS_OBJECT_DETECTOR_FL2
|
|
|
|
for i in tqdm(range(batch_size)):
|
|
img = tensor2pil(images[i]).convert("RGB")
|
|
img_x = rects.rectangles[i][0]
|
|
img_y = rects.rectangles[i][1]
|
|
img_target_width = rects.rectangles[i][2]
|
|
img_target_height = rects.rectangles[i][3]
|
|
|
|
od = LS_OBJECT_DETECTOR_FL2()
|
|
if florence2_model is not None:
|
|
bboxes = od.object_detector_fl2(image=[images[i]], prompt=object_prompt, florence2_model=florence2_model,
|
|
sort_method="confidence", bbox_select="first", select_index="0")[0]
|
|
bbox_mask = self.draw_bbox_mask(img, bboxes, 0, 0, 0, 0)
|
|
resized_img = self.image_auto_crop_v3(img, img_target_width, img_target_height, bbox_mask)
|
|
else:
|
|
resized_img = fit_resize_image(img, img_target_width, img_target_height, fit="crop", resize_sampler=Image.LANCZOS)
|
|
|
|
canvas.paste(resized_img, box=(img_x, img_y))
|
|
canvas.paste(color_image, box=(0, 0), mask=rects_border_image.convert("L"))
|
|
|
|
return (pil2tensor(canvas), 1 - image2mask(rects_border_image),)
|
|
|
|
def draw_bbox_mask(self, image, bboxes, grow_top, grow_bottom, grow_left, grow_right
|
|
):
|
|
|
|
mask = Image.new("L", image.size, color='black')
|
|
for bbox in bboxes:
|
|
try:
|
|
if len(bbox) == 0:
|
|
continue
|
|
else:
|
|
x1, y1, x2, y2 = bbox
|
|
except ValueError:
|
|
if len(bbox) == 0:
|
|
continue
|
|
else:
|
|
x1, y1, x2, y2 = bbox[0]
|
|
w = x2 - x1
|
|
h = y2 - y1
|
|
if grow_top:
|
|
y1 = int(y1 - h * grow_top)
|
|
if grow_bottom:
|
|
y2 = int(y2 + h * grow_bottom)
|
|
if grow_left:
|
|
x1 = int(x1 - w * grow_left)
|
|
if grow_right:
|
|
x2 = int(x2 + w * grow_right)
|
|
if y1 > y2 or x1 > x2:
|
|
continue
|
|
draw = ImageDraw.Draw(mask)
|
|
draw.rectangle([x1, y1, x2, y2], fill='white', outline='white', width=0)
|
|
|
|
return mask
|
|
|
|
def image_auto_crop_v3(self, image, proportional_width, proportional_height, mask,
|
|
):
|
|
|
|
scale_to_length = proportional_width
|
|
_image = image
|
|
ratio = proportional_width / proportional_height
|
|
resize_sampler = Image.LANCZOS
|
|
# calculate target width and height
|
|
if ratio > 1:
|
|
target_width = scale_to_length
|
|
target_height = int(target_width / ratio)
|
|
else:
|
|
target_width = scale_to_length
|
|
target_height = int(target_width / ratio)
|
|
|
|
_mask = mask
|
|
bluredmask = gaussian_blur(_mask, 20).convert('L')
|
|
(mask_x, mask_y, mask_w, mask_h) = mask_area(bluredmask)
|
|
orig_ratio = _image.width / _image.height
|
|
target_ratio = target_width / target_height
|
|
# crop image to target ratio
|
|
if orig_ratio > target_ratio: # crop LiftRight side
|
|
crop_w = int(_image.height * target_ratio)
|
|
crop_h = _image.height
|
|
else: # crop TopBottom side
|
|
crop_w = _image.width
|
|
crop_h = int(_image.width / target_ratio)
|
|
crop_x = mask_w // 2 + mask_x - crop_w // 2
|
|
if crop_x < 0:
|
|
crop_x = 0
|
|
if crop_x + crop_w > _image.width:
|
|
crop_x = _image.width - crop_w
|
|
crop_y = mask_h // 2 + mask_y - crop_h // 2
|
|
if crop_y < 0:
|
|
crop_y = 0
|
|
if crop_y + crop_h > _image.height:
|
|
crop_y = _image.height - crop_h
|
|
crop_image = _image.crop((crop_x, crop_y, crop_x + crop_w, crop_y + crop_h))
|
|
ret_image = crop_image.resize((target_width, target_height), resize_sampler)
|
|
|
|
return ret_image
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayerUtility: Collage": LS_Collage,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayerUtility: Collage": "LayerUtility: Collage(Advance)",
|
|
} |