147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
from __future__ import annotations
|
|
from enum import Enum
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import kornia.filters
|
|
from torch import Tensor
|
|
|
|
|
|
def mask_to_torch(mask: Tensor):
|
|
if len(mask.shape) == 3: # BHW -> B1HW
|
|
mask = mask.unsqueeze(1)
|
|
elif len(mask.shape) == 2: # HW -> B1HW
|
|
mask = mask.unsqueeze(0).unsqueeze(0)
|
|
return mask
|
|
|
|
|
|
def image_to_torch(image: Tensor):
|
|
if len(image.shape) == 3:
|
|
image = image.unsqueeze(0)
|
|
image = image.permute(0, 3, 1, 2) # BHWC -> BCHW
|
|
return image
|
|
|
|
|
|
def to_torch(image: Tensor, mask: Tensor):
|
|
if len(image.shape) == 3:
|
|
image = image.unsqueeze(0)
|
|
image = image.permute(0, 3, 1, 2) # BHWC -> BCHW
|
|
|
|
mask = mask_to_torch(mask)
|
|
if image.shape[2:] != mask.shape[2:]:
|
|
raise ValueError(
|
|
f"Image and mask must be the same size. {image.shape[2:]} != {mask.shape[2:]}"
|
|
)
|
|
return image, mask
|
|
|
|
|
|
def to_comfy(image: Tensor):
|
|
return image.permute(0, 2, 3, 1) # BCHW -> BHWC
|
|
|
|
|
|
def mask_floor(mask: Tensor, threshold: float = 0.99):
|
|
return (mask >= threshold).to(mask.dtype)
|
|
|
|
|
|
# torch pad does not support padding greater than image size with "reflect" mode
|
|
def pad_reflect_once(x: Tensor, original_padding: tuple[int, int, int, int]):
|
|
_, _, h, w = x.shape
|
|
padding = np.array(original_padding)
|
|
size = np.array([w, w, h, h])
|
|
|
|
initial_padding = np.minimum(padding, size - 1)
|
|
additional_padding = padding - initial_padding
|
|
|
|
x = F.pad(x, tuple(initial_padding), mode="reflect")
|
|
if np.any(additional_padding > 0):
|
|
x = F.pad(x, tuple(additional_padding), mode="constant")
|
|
return x
|
|
|
|
|
|
def resize_square(image: Tensor, mask: Tensor, size: int):
|
|
_, _, h, w = image.shape
|
|
pad_w, pad_h, prev_size = 0, 0, w
|
|
if w == size and h == size:
|
|
return image, mask, (pad_w, pad_h, prev_size)
|
|
|
|
if w < h:
|
|
pad_w = h - w
|
|
prev_size = h
|
|
elif h < w:
|
|
pad_h = w - h
|
|
prev_size = w
|
|
image = pad_reflect_once(image, (0, pad_w, 0, pad_h))
|
|
mask = pad_reflect_once(mask, (0, pad_w, 0, pad_h))
|
|
|
|
if image.shape[-1] != size:
|
|
image = F.interpolate(image, size=size, mode="nearest-exact")
|
|
mask = F.interpolate(mask, size=size, mode="nearest-exact")
|
|
|
|
return image, mask, (pad_w, pad_h, prev_size)
|
|
|
|
|
|
def undo_resize_square(image: Tensor, original_size: tuple[int, int, int]):
|
|
_, _, h, w = image.shape
|
|
pad_w, pad_h, prev_size = original_size
|
|
if prev_size != w or prev_size != h:
|
|
image = F.interpolate(image, size=prev_size, mode="bilinear")
|
|
return image[:, :, 0 : prev_size - pad_h, 0 : prev_size - pad_w]
|
|
|
|
|
|
def gaussian_blur(image: Tensor, size: int, sigma: float = 0):
|
|
c = image.shape[-3]
|
|
if sigma <= 0:
|
|
sigma = 0.3 * (size - 1) + 0.8
|
|
return kornia.filters.gaussian_blur2d(image, (size, size), (sigma, sigma))
|
|
|
|
|
|
def binary_erosion(mask: Tensor, radius: int):
|
|
kernel = torch.ones(1, 1, radius * 2 + 1, radius * 2 + 1, device=mask.device)
|
|
mask = F.pad(mask, (radius, radius, radius, radius), mode="constant", value=1)
|
|
mask = F.conv2d(mask, kernel, groups=1)
|
|
mask = (mask == kernel.numel()).to(mask.dtype)
|
|
return mask
|
|
|
|
|
|
def binary_dilation(mask: Tensor, radius: int):
|
|
kernel = torch.ones(1, radius * 2 + 1, device=mask.device)
|
|
mask = kornia.filters.filter2d_separable(mask, kernel, kernel, border_type="constant")
|
|
mask = (mask > 0).to(mask.dtype)
|
|
return mask
|
|
|
|
|
|
class BlurKernel(Enum):
|
|
box = "box"
|
|
linear = "linear"
|
|
gaussian = "gaussian"
|
|
|
|
|
|
def mask_blur(mask: Tensor, size: int, method=BlurKernel.gaussian):
|
|
if size <= 2:
|
|
return mask
|
|
if method is BlurKernel.gaussian:
|
|
return gaussian_blur(mask, size)
|
|
|
|
match method, size:
|
|
case BlurKernel.box, _:
|
|
kernel = torch.ones(1, size, device=mask.device) / size
|
|
case BlurKernel.linear, 3:
|
|
kernel = torch.ones(1, size, device=mask.device) / size
|
|
case BlurKernel.linear, _:
|
|
h = size // 2
|
|
kernel = torch.linspace(1 / h, (h - 1) / h, h, device=mask.device)
|
|
kernel = torch.cat([kernel, torch.ones(1, device=mask.device), kernel.flip(0)])
|
|
kernel = kernel / kernel.sum()
|
|
kernel = kernel.unsqueeze(0)
|
|
case _:
|
|
raise ValueError(f"Unknown blur kernel: {method}")
|
|
|
|
mask = kornia.filters.filter2d_separable(mask, kernel, kernel, border_type="reflect")
|
|
return mask
|
|
|
|
|
|
def make_odd(x):
|
|
if x > 0 and x % 2 == 0:
|
|
return x + 1
|
|
return x
|