720 lines
25 KiB
Python
720 lines
25 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import io
|
|
import math
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageCms
|
|
|
|
from .mask_utils import (
|
|
_BICUBIC,
|
|
_BILINEAR,
|
|
_NEAREST,
|
|
_LANCZOS,
|
|
_BOX,
|
|
bbox_from_mask,
|
|
bbox_from_mask_or_none,
|
|
binary_mask,
|
|
bbox_mask,
|
|
prepare_paste_mask,
|
|
refine_error,
|
|
)
|
|
|
|
|
|
VAE_IMAGE_SIZE = 1024 * 1024
|
|
DEFAULT_RESIZE_METHOD = "lanczos"
|
|
IMAGE_RESIZE_METHODS = ["lanczos", "area", "bicubic", "bilinear", "nearest-exact"]
|
|
SUPPORTED_UPSCALE_METHODS = {"lanczos", "area", "bicubic", "bilinear", "nearest-exact"}
|
|
PREFERRED_KONTEXT_RESOLUTIONS = [
|
|
(672, 1568),
|
|
(688, 1504),
|
|
(720, 1456),
|
|
(752, 1392),
|
|
(800, 1328),
|
|
(832, 1248),
|
|
(880, 1184),
|
|
(944, 1104),
|
|
(1024, 1024),
|
|
(1104, 944),
|
|
(1184, 880),
|
|
(1248, 832),
|
|
(1328, 800),
|
|
(1392, 752),
|
|
(1456, 720),
|
|
(1504, 688),
|
|
(1568, 672),
|
|
]
|
|
|
|
|
|
def calculate_dimensions_rounded(target_area: int, ratio: float, multiple: int) -> tuple[int, int]:
|
|
width = math.sqrt(float(target_area) * float(ratio))
|
|
height = width / float(ratio)
|
|
width = max(multiple, round(width / multiple) * multiple)
|
|
height = max(multiple, round(height / multiple) * multiple)
|
|
return int(width), int(height)
|
|
|
|
|
|
def calculate_dimensions(target_area: int, ratio: float) -> tuple[int, int]:
|
|
return calculate_dimensions_rounded(target_area, ratio, 32)
|
|
|
|
|
|
def calculate_dimensions_area(target_area: int, width: int, height: int, multiple: int | None = None) -> tuple[int, int]:
|
|
scale_by = math.sqrt(float(target_area) / float(width * height))
|
|
target_width = width * scale_by
|
|
target_height = height * scale_by
|
|
if multiple is not None:
|
|
target_width = round(target_width / multiple) * multiple
|
|
target_height = round(target_height / multiple) * multiple
|
|
return max(multiple, int(target_width)), max(multiple, int(target_height))
|
|
return max(1, round(target_width)), max(1, round(target_height))
|
|
|
|
|
|
def flux_kontext_target_size(width: int, height: int) -> tuple[int, int]:
|
|
aspect_ratio = width / height
|
|
_, target_width, target_height = min(
|
|
(abs(aspect_ratio - w / h), w, h)
|
|
for w, h in PREFERRED_KONTEXT_RESOLUTIONS
|
|
)
|
|
return target_width, target_height
|
|
|
|
|
|
def normalize_to_srgb(img: Image.Image | None) -> Image.Image | None:
|
|
if img is None:
|
|
return None
|
|
icc = img.info.get("icc_profile") if hasattr(img, "info") else None
|
|
if icc:
|
|
try:
|
|
src_profile = ImageCms.ImageCmsProfile(io.BytesIO(icc))
|
|
dst_profile = ImageCms.createProfile("sRGB")
|
|
return ImageCms.profileToProfile(
|
|
img,
|
|
src_profile,
|
|
dst_profile,
|
|
outputMode="RGB",
|
|
)
|
|
except Exception:
|
|
return img.convert("RGB")
|
|
return img.convert("RGB")
|
|
|
|
|
|
def composite_masked(
|
|
destination: Image.Image,
|
|
source: Image.Image,
|
|
mask_l: Image.Image,
|
|
*,
|
|
allow_resize: bool = False,
|
|
) -> Image.Image:
|
|
if source.size != destination.size:
|
|
if not allow_resize:
|
|
raise ValueError(
|
|
f"RefineNode composite size mismatch: source size {source.size} "
|
|
f"does not match destination size {destination.size}."
|
|
)
|
|
source = source.resize(destination.size, _BICUBIC)
|
|
if mask_l.size != destination.size:
|
|
if not allow_resize:
|
|
raise ValueError(
|
|
f"RefineNode composite size mismatch: mask size {mask_l.size} "
|
|
f"does not match destination size {destination.size}."
|
|
)
|
|
mask_l = mask_l.resize(destination.size, _BILINEAR)
|
|
|
|
dst = np.asarray(destination.convert("RGB"), dtype=np.float32)
|
|
src = np.asarray(source.convert("RGB"), dtype=np.float32)
|
|
alpha = np.asarray(mask_l.convert("L"), dtype=np.float32) / 255.0
|
|
out = src * alpha[:, :, None] + dst * (1.0 - alpha[:, :, None])
|
|
return Image.fromarray(np.clip(np.rint(out), 0, 255).astype(np.uint8), mode="RGB")
|
|
|
|
|
|
def composite_masked_same_size(
|
|
destination: Image.Image,
|
|
source: Image.Image,
|
|
mask_l: Image.Image,
|
|
) -> Image.Image:
|
|
return composite_masked(destination, source, mask_l, allow_resize=False)
|
|
|
|
|
|
def tensor_image_to_pil(images: torch.Tensor, index: int) -> Image.Image:
|
|
tensor = images.detach().cpu()
|
|
if tensor.ndim == 4:
|
|
tensor = tensor[min(index, tensor.shape[0] - 1)]
|
|
if tensor.ndim != 3:
|
|
raise ValueError(f"Expected IMAGE tensor with 3 or 4 dims, got {tuple(images.shape)}")
|
|
if tensor.shape[-1] == 1:
|
|
tensor = tensor.repeat(1, 1, 3)
|
|
elif tensor.shape[-1] == 4:
|
|
tensor = tensor[..., :3]
|
|
if tensor.shape[-1] != 3:
|
|
raise ValueError(f"Expected IMAGE tensor channel count 1, 3, or 4, got {tensor.shape[-1]}")
|
|
arr = (tensor.clamp(0, 1).numpy() * 255.0 + 0.5).astype(np.uint8)
|
|
return Image.fromarray(arr, mode="RGB")
|
|
|
|
|
|
def pil_to_tensor_image(image: Image.Image) -> torch.Tensor:
|
|
arr = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr)
|
|
|
|
|
|
def image_content_signature(image: Image.Image | None) -> str | None:
|
|
if image is None:
|
|
return None
|
|
thumb = image.convert("RGB").resize((32, 32), _BILINEAR)
|
|
return hashlib.sha1(thumb.tobytes()).hexdigest()
|
|
|
|
|
|
def image_batch_size(images: torch.Tensor | None) -> int:
|
|
if images is None:
|
|
return 0
|
|
if images.ndim == 4:
|
|
return int(images.shape[0])
|
|
return 1
|
|
|
|
|
|
def tensor_image_as_list_item(image: Image.Image) -> torch.Tensor:
|
|
return pil_to_tensor_image(image).unsqueeze(0)
|
|
|
|
|
|
def normalize_axis_angle(angle: float) -> float:
|
|
while angle < -90.0:
|
|
angle += 180.0
|
|
while angle >= 90.0:
|
|
angle -= 180.0
|
|
return angle
|
|
|
|
|
|
def mask_principal_axis_angle(mask_l: Image.Image) -> float | None:
|
|
arr = np.asarray(mask_l.convert("L"), dtype=np.uint8) > 0
|
|
ys, xs = np.nonzero(arr)
|
|
if xs.size < 2:
|
|
return None
|
|
coords = np.stack([xs.astype(np.float64), -ys.astype(np.float64)], axis=1)
|
|
coords -= coords.mean(axis=0, keepdims=True)
|
|
cov = np.cov(coords, rowvar=False)
|
|
if not np.isfinite(cov).all():
|
|
return None
|
|
values, vectors = np.linalg.eigh(cov)
|
|
vector = vectors[:, int(np.argmax(values))]
|
|
if not np.isfinite(vector).all() or np.linalg.norm(vector) <= 1e-8:
|
|
return None
|
|
return normalize_axis_angle(math.degrees(math.atan2(float(vector[1]), float(vector[0]))))
|
|
|
|
|
|
def rotated_product_crop(
|
|
image: Image.Image,
|
|
mask_l: Image.Image,
|
|
angle: float,
|
|
canvas_expand: float = 0.0,
|
|
) -> tuple[Image.Image, Image.Image]:
|
|
rotated_image = image.convert("RGB").rotate(
|
|
angle,
|
|
resample=_BICUBIC,
|
|
expand=True,
|
|
fillcolor=(0, 0, 0),
|
|
)
|
|
rotated_mask = binary_mask(mask_l).rotate(
|
|
angle,
|
|
resample=_NEAREST,
|
|
expand=True,
|
|
fillcolor=0,
|
|
)
|
|
rotated_mask = binary_mask(rotated_mask)
|
|
bbox = bbox_from_mask_or_none(rotated_mask)
|
|
if bbox is None:
|
|
return rotated_image, rotated_mask
|
|
expand = max(0.0, min(1.0, float(canvas_expand)))
|
|
if expand <= 0.0:
|
|
crop_box = bbox
|
|
else:
|
|
canvas_w, canvas_h = rotated_image.size
|
|
x1, y1, x2, y2 = bbox
|
|
crop_box = (
|
|
max(0, int(math.floor(x1 * (1.0 - expand)))),
|
|
max(0, int(math.floor(y1 * (1.0 - expand)))),
|
|
min(canvas_w, int(math.ceil(x2 * (1.0 - expand) + canvas_w * expand))),
|
|
min(canvas_h, int(math.ceil(y2 * (1.0 - expand) + canvas_h * expand))),
|
|
)
|
|
return rotated_image.crop(crop_box), rotated_mask.crop(crop_box)
|
|
|
|
|
|
def rotate_image_canvas(image: Image.Image, angle: float) -> Image.Image:
|
|
return image.convert("RGB").rotate(
|
|
float(angle),
|
|
resample=_BICUBIC,
|
|
expand=True,
|
|
fillcolor=(0, 0, 0),
|
|
)
|
|
|
|
|
|
def stack_image_pils(images: list[Image.Image]) -> torch.Tensor:
|
|
if not images:
|
|
raise ValueError("Cannot output an empty image batch.")
|
|
size = images[0].size
|
|
if any(image.size != size for image in images):
|
|
raise ValueError(
|
|
"RefineNode Match Product Angle batch outputs have different crop sizes. "
|
|
"Split the batch or use inputs with matching rotated product bbox sizes."
|
|
)
|
|
return torch.stack([pil_to_tensor_image(image) for image in images], dim=0)
|
|
|
|
|
|
def upscale_image(samples: torch.Tensor, width: int, height: int, method: str, crop: str = "disabled") -> torch.Tensor:
|
|
import comfy.utils
|
|
|
|
method = (method or DEFAULT_RESIZE_METHOD).strip().lower()
|
|
if method == "nearest":
|
|
method = "nearest-exact"
|
|
if method not in SUPPORTED_UPSCALE_METHODS:
|
|
method = "bicubic"
|
|
return comfy.utils.common_upscale(samples, width, height, method, crop)
|
|
|
|
|
|
def upscale_to_fit_with_padding(samples: torch.Tensor, width: int, height: int, method: str) -> torch.Tensor:
|
|
_, _, old_height, old_width = samples.shape
|
|
scale = min(width / old_width, height / old_height)
|
|
resized_width = max(1, min(width, round(old_width * scale)))
|
|
resized_height = max(1, min(height, round(old_height * scale)))
|
|
resized = upscale_image(samples, resized_width, resized_height, method, "disabled")
|
|
out = samples.new_zeros((samples.shape[0], samples.shape[1], height, width))
|
|
x = (width - resized_width) // 2
|
|
y = (height - resized_height) // 2
|
|
out[:, :, y : y + resized_height, x : x + resized_width] = resized
|
|
return out
|
|
|
|
|
|
def upscale_to_kontext_size(samples: torch.Tensor, width: int, height: int, method: str, crop_mode: str) -> torch.Tensor:
|
|
mode = (crop_mode or "crop").strip().lower()
|
|
if mode == "fill":
|
|
return upscale_to_fit_with_padding(samples, width, height, method)
|
|
crop = "disabled" if mode == "disable" else "center"
|
|
return upscale_image(samples, width, height, method, crop)
|
|
|
|
|
|
def reference_image_transform_metadata(
|
|
source_size: tuple[int, int],
|
|
target_size: tuple[int, int],
|
|
resize_method: str,
|
|
crop_mode: str,
|
|
sizing_mode: str,
|
|
source_signature: str | None = None,
|
|
) -> dict[str, Any]:
|
|
source_width, source_height = source_size
|
|
target_width, target_height = target_size
|
|
mode = (crop_mode or "crop").strip().lower()
|
|
if mode not in {"crop", "disable", "fill"}:
|
|
mode = "crop"
|
|
|
|
source_box = (0, 0, int(source_width), int(source_height))
|
|
target_box = (0, 0, int(target_width), int(target_height))
|
|
|
|
if mode == "crop":
|
|
old_aspect = source_width / source_height
|
|
new_aspect = target_width / target_height
|
|
x = 0
|
|
y = 0
|
|
if old_aspect > new_aspect:
|
|
x = round((source_width - source_width * (new_aspect / old_aspect)) / 2)
|
|
elif old_aspect < new_aspect:
|
|
y = round((source_height - source_height * (old_aspect / new_aspect)) / 2)
|
|
source_box = (
|
|
int(x),
|
|
int(y),
|
|
int(source_width - x),
|
|
int(source_height - y),
|
|
)
|
|
elif mode == "fill":
|
|
scale = min(target_width / source_width, target_height / source_height)
|
|
resized_width = max(1, min(target_width, round(source_width * scale)))
|
|
resized_height = max(1, min(target_height, round(source_height * scale)))
|
|
x = (target_width - resized_width) // 2
|
|
y = (target_height - resized_height) // 2
|
|
target_box = (int(x), int(y), int(x + resized_width), int(y + resized_height))
|
|
|
|
return {
|
|
"source_size": (int(source_width), int(source_height)),
|
|
"target_size": (int(target_width), int(target_height)),
|
|
"source_content_box": source_box,
|
|
"target_content_box": target_box,
|
|
"resize_method": (resize_method or DEFAULT_RESIZE_METHOD).strip().lower(),
|
|
"crop_mode": mode,
|
|
"sizing_mode": sizing_mode,
|
|
"source_signature": source_signature,
|
|
}
|
|
|
|
|
|
def scale_box(
|
|
box: tuple[int, int, int, int],
|
|
from_size: tuple[int, int],
|
|
to_size: tuple[int, int],
|
|
) -> tuple[int, int, int, int]:
|
|
from_width, from_height = from_size
|
|
to_width, to_height = to_size
|
|
if from_width == to_width and from_height == to_height:
|
|
return tuple(int(v) for v in box)
|
|
scale_x = to_width / from_width
|
|
scale_y = to_height / from_height
|
|
x1, y1, x2, y2 = box
|
|
return (
|
|
int(round(x1 * scale_x)),
|
|
int(round(y1 * scale_y)),
|
|
int(round(x2 * scale_x)),
|
|
int(round(y2 * scale_y)),
|
|
)
|
|
|
|
|
|
def clamp_box(box: tuple[int, int, int, int], size: tuple[int, int]) -> tuple[int, int, int, int]:
|
|
width, height = size
|
|
x1, y1, x2, y2 = tuple(int(v) for v in box)
|
|
x1 = max(0, min(width - 1, x1))
|
|
y1 = max(0, min(height - 1, y1))
|
|
x2 = max(x1 + 1, min(width, x2))
|
|
y2 = max(y1 + 1, min(height, y2))
|
|
return x1, y1, x2, y2
|
|
|
|
|
|
def pil_resize_filter(method: str | None) -> int:
|
|
method = (method or DEFAULT_RESIZE_METHOD).strip().lower()
|
|
if method in {"lanczos", "nearest-exact"}:
|
|
return _LANCZOS if method == "lanczos" else _NEAREST
|
|
if method == "area":
|
|
return _BOX
|
|
if method == "bilinear":
|
|
return _BILINEAR
|
|
if method == "nearest":
|
|
return _NEAREST
|
|
return _BICUBIC
|
|
|
|
|
|
def flatten_refine_info_items(info: dict[str, Any] | list[Any] | None) -> list[Any]:
|
|
if info is None:
|
|
return []
|
|
values = info if isinstance(info, list) else [info]
|
|
items = []
|
|
for value in values:
|
|
if isinstance(value, list):
|
|
items.extend(flatten_refine_info_items(value))
|
|
continue
|
|
if not isinstance(value, dict):
|
|
continue
|
|
value_items = value.get("items")
|
|
if isinstance(value_items, list):
|
|
items.extend(value_items)
|
|
elif value_items is not None:
|
|
items.append(value_items)
|
|
return items
|
|
|
|
|
|
def update_info_with_kontext_transforms(
|
|
info: dict[str, Any] | list[Any] | None,
|
|
slot_transforms: dict[str, list[dict[str, Any] | None]],
|
|
) -> dict[str, Any]:
|
|
base_info = {}
|
|
values = info if isinstance(info, list) else [info]
|
|
for value in values:
|
|
if isinstance(value, dict):
|
|
base_info = value.copy()
|
|
break
|
|
items = flatten_refine_info_items(info)
|
|
if not items or not slot_transforms:
|
|
base_info["items"] = items
|
|
return base_info
|
|
|
|
def transform_at(slot: str, index: int) -> dict[str, Any] | None:
|
|
transforms = slot_transforms.get(slot) or []
|
|
if not transforms:
|
|
return None
|
|
return transforms[min(index, len(transforms) - 1)]
|
|
|
|
def item_model_size(item: dict[str, Any]) -> tuple[int, int] | None:
|
|
model_image = item.get("model_image")
|
|
if isinstance(model_image, Image.Image):
|
|
return tuple(int(v) for v in model_image.size)
|
|
size = item.get("model_size") or item.get("model_image_size")
|
|
if isinstance(size, (list, tuple)) and len(size) == 2:
|
|
return (int(size[0]), int(size[1]))
|
|
return None
|
|
|
|
def item_model_signature(item: dict[str, Any]) -> str | None:
|
|
model_image = item.get("model_image")
|
|
if isinstance(model_image, Image.Image):
|
|
return image_content_signature(model_image)
|
|
signature = item.get("model_image_signature")
|
|
if isinstance(signature, str):
|
|
return signature
|
|
return None
|
|
|
|
def choose_transform(item: dict[str, Any], index: int) -> tuple[dict[str, Any] | None, str]:
|
|
model_signature = item_model_signature(item)
|
|
if model_signature is not None:
|
|
for slot in ("image1", "image2", "image3"):
|
|
selected = transform_at(slot, index)
|
|
if selected and selected.get("source_signature") == model_signature:
|
|
return selected, slot
|
|
|
|
model_size = item_model_size(item)
|
|
if model_size is not None:
|
|
for slot in ("image1", "image2", "image3"):
|
|
selected = transform_at(slot, index)
|
|
if selected and tuple(selected.get("source_size", ())) == model_size:
|
|
return selected, slot
|
|
|
|
for slot in ("image1", "image2", "image3"):
|
|
selected = transform_at(slot, index)
|
|
if selected is not None:
|
|
return selected, slot
|
|
return None, "auto"
|
|
|
|
updated_items = []
|
|
for index, item in enumerate(items):
|
|
if not isinstance(item, dict):
|
|
updated_items.append(item)
|
|
continue
|
|
updated = item.copy()
|
|
selected, selected_source = choose_transform(item, index)
|
|
if selected is not None:
|
|
updated["reference_image_transform"] = selected
|
|
updated["reference_image_transform_source"] = selected_source
|
|
updated_items.append(updated)
|
|
|
|
updated_info = base_info.copy()
|
|
updated_info["items"] = updated_items
|
|
return updated_info
|
|
|
|
|
|
def restore_generated_to_model_space(
|
|
generated: Image.Image,
|
|
model_image: Image.Image,
|
|
transform: dict[str, Any] | None,
|
|
) -> Image.Image:
|
|
if not isinstance(transform, dict):
|
|
return generated
|
|
|
|
mode = (transform.get("crop_mode") or "crop").strip().lower()
|
|
target_size = tuple(transform.get("target_size") or generated.size)
|
|
if len(target_size) != 2 or target_size[0] <= 0 or target_size[1] <= 0:
|
|
target_size = generated.size
|
|
|
|
if mode == "fill":
|
|
target_box = tuple(transform.get("target_content_box") or (0, 0, target_size[0], target_size[1]))
|
|
target_box = scale_box(target_box, target_size, generated.size)
|
|
target_box = clamp_box(target_box, generated.size)
|
|
content = generated.crop(target_box)
|
|
return content.resize(model_image.size, _BICUBIC)
|
|
|
|
if mode == "crop":
|
|
source_box = tuple(transform.get("source_content_box") or (0, 0, *model_image.size))
|
|
source_box = clamp_box(source_box, model_image.size)
|
|
x1, y1, x2, y2 = source_box
|
|
restored = model_image.copy()
|
|
content = generated.resize((x2 - x1, y2 - y1), _BICUBIC)
|
|
restored.paste(content, (x1, y1))
|
|
return restored
|
|
|
|
return generated.resize(model_image.size, _BICUBIC)
|
|
|
|
|
|
def restore_mask_to_model_space(
|
|
mask_l: Image.Image,
|
|
model_size: tuple[int, int],
|
|
transform: dict[str, Any] | None,
|
|
) -> Image.Image:
|
|
mask_l = binary_mask(mask_l)
|
|
if not isinstance(transform, dict):
|
|
return mask_l.resize(model_size, _NEAREST)
|
|
|
|
mode = (transform.get("crop_mode") or "crop").strip().lower()
|
|
target_size = tuple(transform.get("target_size") or mask_l.size)
|
|
if len(target_size) != 2 or target_size[0] <= 0 or target_size[1] <= 0:
|
|
target_size = mask_l.size
|
|
|
|
if mask_l.size != tuple(target_size):
|
|
mask_l = mask_l.resize(tuple(target_size), _NEAREST)
|
|
|
|
if mode == "fill":
|
|
target_box = tuple(transform.get("target_content_box") or (0, 0, target_size[0], target_size[1]))
|
|
target_box = clamp_box(target_box, tuple(target_size))
|
|
content = mask_l.crop(target_box)
|
|
return binary_mask(content.resize(model_size, _NEAREST))
|
|
|
|
if mode == "crop":
|
|
source_box = tuple(transform.get("source_content_box") or (0, 0, *model_size))
|
|
source_box = clamp_box(source_box, model_size)
|
|
restored = Image.new("L", model_size, 0)
|
|
x1, y1, x2, y2 = source_box
|
|
content = mask_l.resize((x2 - x1, y2 - y1), _NEAREST)
|
|
restored.paste(content, (x1, y1))
|
|
return binary_mask(restored)
|
|
|
|
return binary_mask(mask_l.resize(model_size, _NEAREST))
|
|
|
|
|
|
def safe_size_tuple(value: Any, fallback: tuple[int, int]) -> tuple[int, int]:
|
|
if isinstance(value, (list, tuple)) and len(value) == 2:
|
|
try:
|
|
width = int(value[0])
|
|
height = int(value[1])
|
|
except (TypeError, ValueError):
|
|
return fallback
|
|
if width > 0 and height > 0:
|
|
return width, height
|
|
return fallback
|
|
|
|
|
|
def transformed_content_box(
|
|
transform: dict[str, Any],
|
|
key: str,
|
|
from_size: tuple[int, int],
|
|
to_size: tuple[int, int],
|
|
fallback: tuple[int, int, int, int],
|
|
) -> tuple[int, int, int, int]:
|
|
box = tuple(transform.get(key) or fallback)
|
|
if len(box) != 4:
|
|
box = fallback
|
|
box = tuple(int(value) for value in box)
|
|
if from_size != to_size:
|
|
box = scale_box(box, from_size, to_size)
|
|
return clamp_box(box, to_size)
|
|
|
|
|
|
def model_image_to_generated_space(
|
|
model_image: Image.Image,
|
|
generated_size: tuple[int, int],
|
|
transform: dict[str, Any] | None,
|
|
) -> Image.Image:
|
|
image = model_image.convert("RGB")
|
|
if generated_size[0] <= 0 or generated_size[1] <= 0:
|
|
raise ValueError("Generated image size must be positive.")
|
|
if not isinstance(transform, dict):
|
|
return image.resize(generated_size, _BICUBIC)
|
|
|
|
mode = (transform.get("crop_mode") or "crop").strip().lower()
|
|
target_size = safe_size_tuple(transform.get("target_size"), generated_size)
|
|
source_size = safe_size_tuple(transform.get("source_size"), image.size)
|
|
resample = pil_resize_filter(transform.get("resize_method"))
|
|
|
|
if mode == "fill":
|
|
target_box = transformed_content_box(
|
|
transform,
|
|
"target_content_box",
|
|
target_size,
|
|
generated_size,
|
|
(0, 0, target_size[0], target_size[1]),
|
|
)
|
|
canvas = Image.new("RGB", generated_size, (0, 0, 0))
|
|
x1, y1, x2, y2 = target_box
|
|
content = image.resize((x2 - x1, y2 - y1), resample)
|
|
canvas.paste(content, (x1, y1))
|
|
return canvas
|
|
|
|
if mode == "crop":
|
|
source_box = transformed_content_box(
|
|
transform,
|
|
"source_content_box",
|
|
source_size,
|
|
image.size,
|
|
(0, 0, source_size[0], source_size[1]),
|
|
)
|
|
return image.crop(source_box).resize(generated_size, resample)
|
|
|
|
return image.resize(generated_size, resample)
|
|
|
|
|
|
def model_mask_to_generated_space(
|
|
mask_l: Image.Image,
|
|
generated_size: tuple[int, int],
|
|
transform: dict[str, Any] | None,
|
|
) -> Image.Image:
|
|
mask_l = binary_mask(mask_l)
|
|
if generated_size[0] <= 0 or generated_size[1] <= 0:
|
|
raise ValueError("Generated image size must be positive.")
|
|
if not isinstance(transform, dict):
|
|
return binary_mask(mask_l.resize(generated_size, _NEAREST))
|
|
|
|
mode = (transform.get("crop_mode") or "crop").strip().lower()
|
|
target_size = safe_size_tuple(transform.get("target_size"), generated_size)
|
|
source_size = safe_size_tuple(transform.get("source_size"), mask_l.size)
|
|
|
|
if mode == "fill":
|
|
target_box = transformed_content_box(
|
|
transform,
|
|
"target_content_box",
|
|
target_size,
|
|
generated_size,
|
|
(0, 0, target_size[0], target_size[1]),
|
|
)
|
|
canvas = Image.new("L", generated_size, 0)
|
|
x1, y1, x2, y2 = target_box
|
|
content = mask_l.resize((x2 - x1, y2 - y1), _NEAREST)
|
|
canvas.paste(content, (x1, y1))
|
|
return binary_mask(canvas)
|
|
|
|
if mode == "crop":
|
|
source_box = transformed_content_box(
|
|
transform,
|
|
"source_content_box",
|
|
source_size,
|
|
mask_l.size,
|
|
(0, 0, source_size[0], source_size[1]),
|
|
)
|
|
return binary_mask(mask_l.crop(source_box).resize(generated_size, _NEAREST))
|
|
|
|
return binary_mask(mask_l.resize(generated_size, _NEAREST))
|
|
|
|
|
|
def model_paste_mask_to_generated_space(
|
|
mask_l: Image.Image,
|
|
generated_size: tuple[int, int],
|
|
transform: dict[str, Any] | None,
|
|
) -> Image.Image:
|
|
mask_l = mask_l.convert("L")
|
|
if generated_size[0] <= 0 or generated_size[1] <= 0:
|
|
raise ValueError("Generated image size must be positive.")
|
|
if not isinstance(transform, dict):
|
|
return mask_l.resize(generated_size, _BILINEAR)
|
|
|
|
mode = (transform.get("crop_mode") or "crop").strip().lower()
|
|
target_size = safe_size_tuple(transform.get("target_size"), generated_size)
|
|
source_size = safe_size_tuple(transform.get("source_size"), mask_l.size)
|
|
|
|
if mode == "fill":
|
|
target_box = transformed_content_box(
|
|
transform,
|
|
"target_content_box",
|
|
target_size,
|
|
generated_size,
|
|
(0, 0, target_size[0], target_size[1]),
|
|
)
|
|
canvas = Image.new("L", generated_size, 0)
|
|
x1, y1, x2, y2 = target_box
|
|
content = mask_l.resize((x2 - x1, y2 - y1), _BILINEAR)
|
|
canvas.paste(content, (x1, y1))
|
|
return canvas
|
|
|
|
if mode == "crop":
|
|
source_box = transformed_content_box(
|
|
transform,
|
|
"source_content_box",
|
|
source_size,
|
|
mask_l.size,
|
|
(0, 0, source_size[0], source_size[1]),
|
|
)
|
|
return mask_l.crop(source_box).resize(generated_size, _BILINEAR)
|
|
|
|
return mask_l.resize(generated_size, _BILINEAR)
|
|
|
|
|
|
def restore_mask_to_original_space(mask_l: Image.Image, item: dict[str, Any]) -> Image.Image:
|
|
original = item["origin_image"]
|
|
model_image = item["model_image"]
|
|
model_size = model_image.size if isinstance(model_image, Image.Image) else tuple(item.get("model_size") or original.size)
|
|
model_mask = restore_mask_to_model_space(mask_l, model_size, item.get("reference_image_transform"))
|
|
crop_box = item.get("crop_box")
|
|
if crop_box:
|
|
crop_box = tuple(int(value) for value in crop_box)
|
|
restored = Image.new("L", original.size, 0)
|
|
restored.paste(model_mask.resize((crop_box[2] - crop_box[0], crop_box[3] - crop_box[1]), _NEAREST), (crop_box[0], crop_box[1]))
|
|
return binary_mask(restored)
|
|
return binary_mask(model_mask.resize(original.size, _NEAREST))
|