475 lines
16 KiB
Python
475 lines
16 KiB
Python
import torch
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
from diffusers import FluxFillPipeline, FluxPriorReduxPipeline
|
|
import math
|
|
import torch
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
|
|
def standardize_and_binarize_mask(mask_input):
|
|
"""
|
|
Converts a ComfyUI MASK input (torch.Tensor) to a standardized
|
|
2D NumPy array (H, W), uint8, binarized (0 or 255).
|
|
"""
|
|
if isinstance(mask_input, torch.Tensor):
|
|
mask_np = mask_input.cpu().numpy()
|
|
elif isinstance(mask_input, np.ndarray):
|
|
mask_np = mask_input
|
|
else:
|
|
raise TypeError(f"Mask input type {type(mask_input)} not supported. Expected torch.Tensor or np.ndarray.")
|
|
|
|
# Handle common ComfyUI mask shapes like (N, H, W) or (N, C, H, W)
|
|
if mask_np.ndim == 3 and mask_np.shape[0] == 1: # (1, H, W)
|
|
mask_np = mask_np.squeeze(0)
|
|
elif mask_np.ndim == 4 and mask_np.shape[0] == 1 and mask_np.shape[1] == 1: # (1,1,H,W)
|
|
mask_np = mask_np.squeeze(0).squeeze(0)
|
|
|
|
|
|
# Ensure single channel (H, W)
|
|
if mask_np.ndim == 3:
|
|
if mask_np.shape[-1] == 1: # (H, W, 1)
|
|
mask_np = mask_np.squeeze(axis=-1)
|
|
elif mask_np.shape[0] == 1: # Also (1, H, W) if not caught above
|
|
mask_np = mask_np.squeeze(0)
|
|
elif mask_np.shape[-1] > 1: # (H, W, C) with C > 1, e.g., RGB mask passed incorrectly
|
|
|
|
mask_np = mask_np[..., 0]
|
|
|
|
|
|
|
|
if mask_np.ndim != 2:
|
|
raise ValueError(f"Mask has unexpected shape {mask_np.shape} after initial processing. Expected 2D mask (H,W). Original input shape was {mask_input.shape if hasattr(mask_input, 'shape') else 'unknown'}.")
|
|
|
|
# Normalize to 0-255 uint8 if it's a float mask (0.0-1.0)
|
|
if mask_np.dtype == np.float32 or mask_np.dtype == np.float64:
|
|
if mask_np.min() >= 0.0 and mask_np.max() <= 1.0:
|
|
mask_np = (mask_np * 255).astype(np.uint8)
|
|
else: # Floats not in 0-1 range, clip and convert
|
|
mask_np = np.clip(mask_np, 0, 255).astype(np.uint8)
|
|
elif mask_np.dtype != np.uint8: # Other types, attempt to convert
|
|
mask_np = np.clip(mask_np, 0, 255).astype(np.uint8)
|
|
|
|
# Binarize the mask: values > 128 become 255, others 0.
|
|
_, mask_np_binarized = cv2.threshold(mask_np, 128, 255, cv2.THRESH_BINARY)
|
|
return mask_np_binarized
|
|
|
|
|
|
def create_highlighted_mask(image_np, mask_np, alpha=0.5, gray_value=128):
|
|
|
|
|
|
if mask_np.max() <= 1.0:
|
|
mask_np = (mask_np * 255).astype(np.uint8)
|
|
mask_bool = mask_np > 128
|
|
|
|
image_float = image_np.astype(np.float32)
|
|
|
|
gray_overlay = np.full_like(image_float, gray_value, dtype=np.float32)
|
|
|
|
result = image_float.copy()
|
|
result[mask_bool] = (
|
|
(1 - alpha) * image_float[mask_bool] + alpha * gray_overlay[mask_bool]
|
|
)
|
|
|
|
return result.astype(np.uint8)
|
|
|
|
def f(r, T=0.6, beta=0.1):
|
|
return np.where(r < T, beta + (1 - beta) / T * r, 1)
|
|
|
|
# Get the bounding box of the mask
|
|
def get_bbox_from_mask(mask):
|
|
h,w = mask.shape[0],mask.shape[1]
|
|
|
|
if mask.sum() < 10:
|
|
return 0,h,0,w
|
|
rows = np.any(mask,axis=1)
|
|
cols = np.any(mask,axis=0)
|
|
y1,y2 = np.where(rows)[0][[0,-1]]
|
|
x1,x2 = np.where(cols)[0][[0,-1]]
|
|
return (y1,y2,x1,x2)
|
|
|
|
# Expand the bounding box
|
|
def expand_bbox(mask, yyxx, ratio, min_crop=0):
|
|
y1,y2,x1,x2 = yyxx
|
|
H,W = mask.shape[0], mask.shape[1]
|
|
|
|
yyxx_area = (y2-y1+1) * (x2-x1+1)
|
|
r1 = yyxx_area / (H * W)
|
|
r2 = f(r1)
|
|
ratio = math.sqrt(r2 / r1)
|
|
|
|
xc, yc = 0.5 * (x1 + x2), 0.5 * (y1 + y2)
|
|
h = ratio * (y2-y1+1)
|
|
w = ratio * (x2-x1+1)
|
|
h = max(h,min_crop)
|
|
w = max(w,min_crop)
|
|
|
|
x1 = int(xc - w * 0.5)
|
|
x2 = int(xc + w * 0.5)
|
|
y1 = int(yc - h * 0.5)
|
|
y2 = int(yc + h * 0.5)
|
|
|
|
x1 = max(0,x1)
|
|
x2 = min(W,x2)
|
|
y1 = max(0,y1)
|
|
y2 = min(H,y2)
|
|
return (y1,y2,x1,x2)
|
|
|
|
# Pad the image to a square shape
|
|
def pad_to_square(image, pad_value = 255, random = False):
|
|
H,W = image.shape[0], image.shape[1]
|
|
if H == W:
|
|
return image
|
|
|
|
padd = abs(H - W)
|
|
if random:
|
|
padd_1 = int(np.random.randint(0,padd))
|
|
else:
|
|
padd_1 = int(padd / 2)
|
|
padd_2 = padd - padd_1
|
|
|
|
if len(image.shape) == 2:
|
|
if H > W:
|
|
pad_param = ((0, 0), (padd_1, padd_2))
|
|
else:
|
|
pad_param = ((padd_1, padd_2), (0, 0))
|
|
elif len(image.shape) == 3:
|
|
if H > W:
|
|
pad_param = ((0, 0), (padd_1, padd_2), (0, 0))
|
|
else:
|
|
pad_param = ((padd_1, padd_2), (0, 0), (0, 0))
|
|
|
|
image = np.pad(image, pad_param, 'constant', constant_values=pad_value)
|
|
|
|
return image
|
|
|
|
# Expand the image and mask
|
|
def expand_image_mask(image, mask, ratio=1.4):
|
|
h,w = image.shape[0], image.shape[1]
|
|
H,W = int(h * ratio), int(w * ratio)
|
|
h1 = int((H - h) // 2)
|
|
h2 = H - h - h1
|
|
w1 = int((W -w) // 2)
|
|
w2 = W -w - w1
|
|
|
|
pad_param_image = ((h1,h2),(w1,w2),(0,0))
|
|
pad_param_mask = ((h1,h2),(w1,w2))
|
|
image = np.pad(image, pad_param_image, 'constant', constant_values=255)
|
|
mask = np.pad(mask, pad_param_mask, 'constant', constant_values=0)
|
|
return image, mask
|
|
|
|
# Convert the bounding box to a square shape
|
|
def box2squre(image, box):
|
|
H,W = image.shape[0], image.shape[1]
|
|
y1,y2,x1,x2 = box
|
|
cx = (x1 + x2) // 2
|
|
cy = (y1 + y2) // 2
|
|
h,w = y2-y1, x2-x1
|
|
|
|
if h >= w:
|
|
x1 = cx - h//2
|
|
x2 = cx + h//2
|
|
else:
|
|
y1 = cy - w//2
|
|
y2 = cy + w//2
|
|
x1 = max(0,x1)
|
|
x2 = min(W,x2)
|
|
y1 = max(0,y1)
|
|
y2 = min(H,y2)
|
|
return (y1,y2,x1,x2)
|
|
|
|
# Crop the predicted image back to the original image
|
|
def crop_back( pred, tar_image, extra_sizes, tar_box_yyxx_crop):
|
|
H1, W1, H2, W2 = extra_sizes
|
|
y1,y2,x1,x2 = tar_box_yyxx_crop
|
|
pred = cv2.resize(pred, (W2, H2))
|
|
m = 2 # maigin_pixel
|
|
|
|
if W1 == H1:
|
|
if m != 0:
|
|
tar_image[y1+m :y2-m, x1+m:x2-m, :] = pred[m:-m, m:-m]
|
|
else:
|
|
tar_image[y1 :y2, x1:x2, :] = pred[:, :]
|
|
return tar_image
|
|
|
|
if W1 < W2:
|
|
pad1 = int((W2 - W1) / 2)
|
|
pad2 = W2 - W1 - pad1
|
|
pred = pred[:,pad1: -pad2, :]
|
|
else:
|
|
pad1 = int((H2 - H1) / 2)
|
|
pad2 = H2 - H1 - pad1
|
|
pred = pred[pad1: -pad2, :, :]
|
|
|
|
gen_image = tar_image.copy()
|
|
if m != 0:
|
|
gen_image[y1+m :y2-m, x1+m:x2-m, :] = pred[m:-m, m:-m]
|
|
else:
|
|
gen_image[y1 :y2, x1:x2, :] = pred[:, :]
|
|
|
|
return gen_image
|
|
|
|
|
|
class MaskOption:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
|
|
"sketch": ("MASK", ),
|
|
"upload": ("MASK", ),
|
|
"option": (["sketch", "upload"], {"default": "sketch"}),
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", )
|
|
FUNCTION = "MaskOption"
|
|
|
|
def MaskOption(self, sketch, upload, option):
|
|
|
|
if option == "sketch":
|
|
mask = sketch
|
|
elif option == "upload":
|
|
mask = upload
|
|
|
|
return mask
|
|
|
|
|
|
|
|
class ReduxProcess:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"ref_image": ("IMAGE", ),
|
|
"ref_mask": ("MASK", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
FUNCTION = "ReduxProcess"
|
|
|
|
def ReduxProcess(self, ref_image, ref_mask):
|
|
|
|
ref_image_np = np.array(ref_image)[0] # Assuming ref_image is a batch of 1
|
|
ref_image_np = (ref_image_np * 255).round().astype(np.uint8)
|
|
|
|
# Standardize and binarize the reference mask
|
|
ref_mask_np = standardize_and_binarize_mask(ref_mask) # ref_mask is the MASK input
|
|
|
|
# Remove the background information of the reference picture
|
|
# Ensure ref_mask_np is 2D before passing to get_bbox_from_mask
|
|
ref_box_yyxx = get_bbox_from_mask(ref_mask_np)
|
|
|
|
# Create a 0/1 mask for multiplication (0 for background, 1 for foreground)
|
|
ref_mask_0_1 = ref_mask_np // 255
|
|
ref_mask_3_0_1 = np.stack([ref_mask_0_1, ref_mask_0_1, ref_mask_0_1], axis=-1)
|
|
|
|
# Apply mask: foreground where mask is 1, white background where mask is 0
|
|
masked_ref_image = ref_image_np * ref_mask_3_0_1 + \
|
|
np.ones_like(ref_image_np, dtype=np.uint8) * 255 * (1 - ref_mask_3_0_1)
|
|
masked_ref_image = masked_ref_image.astype(np.uint8)
|
|
|
|
# Extract the box where the reference image is located, and place the reference object at the center of the image
|
|
y1,y2,x1,x2 = ref_box_yyxx
|
|
# Ensure coordinates are valid for slicing
|
|
y1, y2 = max(0, y1), min(masked_ref_image.shape[0], y2)
|
|
x1, x2 = max(0, x1), min(masked_ref_image.shape[1], x2)
|
|
|
|
if y1 >= y2 or x1 >= x2: # Handle empty crop case
|
|
|
|
pass
|
|
|
|
masked_ref_image = masked_ref_image[y1:y2, x1:x2, :]
|
|
ref_mask_processed = ref_mask_np[y1:y2, x1:x2] # Use the binarized mask for consistency
|
|
ratio = 1.3
|
|
masked_ref_image, ref_mask_expanded = expand_image_mask(masked_ref_image, ref_mask_processed, ratio=ratio)
|
|
masked_ref_image = pad_to_square(masked_ref_image, pad_value = 255, random = False)
|
|
# ref_mask_expanded is the corresponding mask for the expanded image, if needed later.
|
|
|
|
|
|
# Extract the features of the reference image
|
|
masked_ref_image = cv2.resize(masked_ref_image.astype(np.uint8), (768, 768)).astype(np.uint8)
|
|
|
|
masked_ref_image = torch.from_numpy(masked_ref_image).unsqueeze(0).float() / 255.0
|
|
|
|
return (masked_ref_image, )
|
|
|
|
|
|
|
|
class FillProcess:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source_image": ("IMAGE", ),
|
|
"ref_image": ("IMAGE", ),
|
|
"source_mask": ("MASK", ),
|
|
"ref_mask": ("MASK", ),
|
|
"iterations": ("INT", {"default": 2}),
|
|
|
|
}
|
|
}
|
|
RETURN_TYPES = (
|
|
"IMAGE", "MASK", "H1", "W1", "H2", "W2", "old_tar_image", "tar_box_yyxx_crop", "IMAGE",
|
|
)
|
|
FUNCTION = "FillProcess"
|
|
|
|
def FillProcess(self, source_image, ref_image, source_mask, ref_mask, iterations):
|
|
|
|
size = (768, 768)
|
|
|
|
ref_image_np = np.array(ref_image)[0] # Assuming batch of 1
|
|
tar_image_np = np.array(source_image)[0] # Assuming batch of 1
|
|
|
|
ref_image_np = (ref_image_np * 255).round().astype(np.uint8)
|
|
tar_image_np = (tar_image_np * 255).round().astype(np.uint8)
|
|
|
|
# Standardize and binarize masks
|
|
ref_mask_np = standardize_and_binarize_mask(ref_mask) # ref_mask is MASK input
|
|
tar_mask_np = standardize_and_binarize_mask(source_mask) # source_mask is MASK input
|
|
|
|
# --- Reference Image Processing (similar to ReduxProcess) ---
|
|
ref_box_yyxx = get_bbox_from_mask(ref_mask_np)
|
|
|
|
ref_mask_0_1 = ref_mask_np // 255
|
|
ref_mask_3_0_1 = np.stack([ref_mask_0_1, ref_mask_0_1, ref_mask_0_1], axis=-1)
|
|
|
|
masked_ref_image = ref_image_np * ref_mask_3_0_1 + \
|
|
np.ones_like(ref_image_np, dtype=np.uint8) * 255 * (1 - ref_mask_3_0_1)
|
|
masked_ref_image = masked_ref_image.astype(np.uint8)
|
|
|
|
y1_ref, y2_ref, x1_ref, x2_ref = ref_box_yyxx
|
|
y1_ref, y2_ref = max(0, y1_ref), min(masked_ref_image.shape[0], y2_ref)
|
|
x1_ref, x2_ref = max(0, x1_ref), min(masked_ref_image.shape[1], x2_ref)
|
|
|
|
if not(y1_ref >= y2_ref or x1_ref >= x2_ref): # Check for valid crop
|
|
masked_ref_image = masked_ref_image[y1_ref:y2_ref, x1_ref:x2_ref, :]
|
|
ref_mask_cropped = ref_mask_np[y1_ref:y2_ref, x1_ref:x2_ref]
|
|
else: # Fallback if bbox is invalid
|
|
ref_mask_cropped = ref_mask_np # use original mask if crop failed
|
|
|
|
ratio_ref = 1.3
|
|
masked_ref_image, _ = expand_image_mask(masked_ref_image, ref_mask_cropped, ratio=ratio_ref) # We mainly need the image
|
|
masked_ref_image = pad_to_square(masked_ref_image, pad_value = 255, random = False)
|
|
# --- End Reference Image Processing ---
|
|
|
|
# Dilate the target mask
|
|
kernel = np.ones((7, 7), np.uint8)
|
|
tar_mask_dilated = cv2.dilate(tar_mask_np, kernel, iterations=iterations)
|
|
|
|
# Zoom in on the target image based on the dilated target mask
|
|
tar_box_yyxx = get_bbox_from_mask(tar_mask_dilated)
|
|
tar_box_yyxx_expanded = expand_bbox(tar_mask_dilated, tar_box_yyxx, ratio=1.2) # Pass tar_mask_dilated
|
|
|
|
|
|
tar_box_yyxx_crop = expand_bbox(tar_image_np, tar_box_yyxx_expanded, ratio=2)
|
|
tar_box_yyxx_crop = expand_bbox(tar_mask_dilated, tar_box_yyxx_expanded, ratio=2)
|
|
tar_box_yyxx_crop = box2squre(tar_image_np, tar_box_yyxx_crop) # box2square uses image for H,W
|
|
|
|
y1_tar, y2_tar, x1_tar, x2_tar = tar_box_yyxx_crop
|
|
|
|
old_tar_image_np = tar_image_np.copy()
|
|
|
|
# Ensure crop coordinates are valid
|
|
y1_tar, y2_tar = max(0, y1_tar), min(tar_image_np.shape[0], y2_tar)
|
|
x1_tar, x2_tar = max(0, x1_tar), min(tar_image_np.shape[1], x2_tar)
|
|
|
|
if y1_tar >= y2_tar or x1_tar >= x2_tar: # Handle empty crop
|
|
# Fallback: use full image, or handle error
|
|
# This indicates an issue with bbox calculation, possibly empty mask
|
|
processed_tar_image = tar_image_np.copy()
|
|
processed_tar_mask = tar_mask_dilated.copy() # or tar_mask_np
|
|
else:
|
|
processed_tar_image = tar_image_np[y1_tar:y2_tar, x1_tar:x2_tar, :]
|
|
processed_tar_mask = tar_mask_dilated[y1_tar:y2_tar, x1_tar:x2_tar]
|
|
|
|
|
|
H1, W1 = processed_tar_image.shape[0], processed_tar_image.shape[1]
|
|
|
|
processed_tar_mask = pad_to_square(processed_tar_mask, pad_value=0)
|
|
processed_tar_mask = cv2.resize(processed_tar_mask, size) # size is (768,768)
|
|
|
|
# Resize reference image (already processed and padded to square)
|
|
masked_ref_image_resized = cv2.resize(masked_ref_image.astype(np.uint8), size).astype(np.uint8)
|
|
|
|
processed_tar_image_padded = pad_to_square(processed_tar_image, pad_value=255)
|
|
H2, W2 = processed_tar_image_padded.shape[0], processed_tar_image_padded.shape[1]
|
|
|
|
processed_tar_image_resized = cv2.resize(processed_tar_image_padded, size)
|
|
|
|
diptych_ref_tar = np.concatenate([masked_ref_image_resized, processed_tar_image_resized], axis=1)
|
|
|
|
# Create diptych mask (mask is for the target part)
|
|
processed_tar_mask_3ch = np.stack([processed_tar_mask, processed_tar_mask, processed_tar_mask], -1)
|
|
|
|
mask_black_ref_part = np.zeros_like(masked_ref_image_resized, dtype=np.uint8) # Black mask for ref part
|
|
mask_diptych = np.concatenate([mask_black_ref_part, processed_tar_mask_3ch], axis=1)
|
|
|
|
show_diptych_ref_tar = create_highlighted_mask(diptych_ref_tar, mask_diptych)
|
|
|
|
|
|
output_mask_diptych_np = mask_diptych[:, :, 0]
|
|
|
|
diptych_ref_tar_tensor = torch.from_numpy(diptych_ref_tar.copy()).unsqueeze(0).float() / 255.0
|
|
mask_diptych_tensor = torch.from_numpy(output_mask_diptych_np.copy()).unsqueeze(0).unsqueeze(1).float() / 255.0 # (B, C, H, W) so (1,1,H,W)
|
|
|
|
mask_diptych_tensor = torch.from_numpy(output_mask_diptych_np.copy()).unsqueeze(0).float() / 255.0
|
|
|
|
|
|
|
|
show_diptych_ref_tar_tensor = torch.from_numpy(show_diptych_ref_tar.copy()).unsqueeze(0).float() / 255.0
|
|
|
|
return (diptych_ref_tar_tensor, mask_diptych_tensor, H1, W1, H2, W2, old_tar_image_np, tar_box_yyxx_crop, show_diptych_ref_tar_tensor)
|
|
|
|
|
|
|
|
|
|
class CropBack:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"raw_image": ("IMAGE",),
|
|
"H1": ("H1",),
|
|
"W1": ("W1",),
|
|
"H2": ("H2",),
|
|
"W2": ("W2",),
|
|
"old_tar_image": ("old_tar_image",),
|
|
"tar_box_yyxx_crop": ("tar_box_yyxx_crop",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "CropBack"
|
|
|
|
def CropBack(
|
|
self,
|
|
raw_image,
|
|
H1, W1, H2, W2,
|
|
old_tar_image,
|
|
tar_box_yyxx_crop
|
|
):
|
|
|
|
raw_image = np.array(raw_image)[0]
|
|
raw_image = (raw_image * 255).round().astype(np.uint8)
|
|
raw_image = Image.fromarray(raw_image)
|
|
|
|
width, height = raw_image.size
|
|
left = width // 2
|
|
right = width
|
|
top = 0
|
|
bottom = height
|
|
edited_image = raw_image.crop((left, top, right, bottom))
|
|
|
|
edited_image = np.array(edited_image)
|
|
edited_image = crop_back(edited_image, old_tar_image, np.array([H1, W1, H2, W2]), np.array(tar_box_yyxx_crop))
|
|
|
|
edited_image = torch.from_numpy(edited_image).unsqueeze(0).float() / 255.0
|
|
|
|
return (edited_image,)
|