util.py

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