univa/utils/get_mask.py

from typing import List
import numpy as np
import cv2
from PIL import Image, ImageChops
import torch
import torch.nn.functional as F

def concat_images_row(images: List[Image.Image], bg_color=(0, 0, 0)) -> Image.Image:
    """
    将多张 PIL 图像按行（横向）拼接在一起。
    
    Args:
        images: 要拼接的图像列表。
        bg_color: 背景颜色，默认黑色；如果图像有透明通道，可以用 (0,0,0,0)。
    
    Returns:
        一张横向拼接后的新图。
    """
    if not images:
        raise ValueError("images 列表不能为空")
    
    # 统一 mode，如果有透明通道，则用 RGBA，否则用 RGB
    modes = {img.mode for img in images}
    mode = "RGBA" if any(m.endswith("A") for m in modes) else "RGB"

    # 计算拼接后画布的尺寸
    widths, heights = zip(*(img.size for img in images))
    total_width = sum(widths)
    max_height = max(heights)

    # 新建画布
    canvas = Image.new(mode, (total_width, max_height), bg_color)

    # 依次粘贴
    x_offset = 0
    for img in images:
        # 如果 img 的 mode 不同，先转换
        if img.mode != mode:
            img = img.convert(mode)
        canvas.paste(img, (x_offset, 0), img if mode=="RGBA" else None)
        x_offset += img.width

    return canvas


def downsample_mask_pytorch(pil_mask: Image.Image, factor: int) -> Image.Image:
    """
    用 PyTorch 的 max_pool2d 对二值 mask 进行下采样，保留块内任何白色。
    
    Args:
        pil_mask: mode='1' 或 'L' 的二值 PIL 图（0/255）。
        factor: 下采样倍数（stride 和 kernel 大小都设为这个值）。
    
    Returns:
        下采样后的二值 PIL Image（mode='1'）。
    """
    # 转成 0/1 float tensor，形状 [1,1,H,W]
    arr = np.array(pil_mask.convert('L'), dtype=np.uint8)
    tensor = torch.from_numpy(arr).float().div_(255.0).unsqueeze(0).unsqueeze(0)
    
    # 用 max_pool2d，下采样
    pooled = F.max_pool2d(tensor, kernel_size=factor, stride=factor)
    
    # 恢复成 0/255，并转回 PIL
    out = (pooled.squeeze(0).squeeze(0) > 0).to(torch.uint8).mul_(255).cpu().numpy()
    return Image.fromarray(out, mode='L').convert('1')

def create_all_white_like(pil_img: Image.Image) -> Image.Image:
    """
    给定一个 PIL 图像，返回一张同样大小的全白二值图（mode='1'）。
    """
    w, h = pil_img.size
    white_array = np.ones((h, w), dtype=np.uint8) * 255  # 注意 shape 是 (H, W)
    return Image.fromarray(white_array, mode='L').convert('1')

def union_masks_np(masks: List[Image.Image]) -> Image.Image:
    """
    接受一个 PIL.Image 列表（mode='1' 或 'L' 的二值图），
    返回它们的并集。
    """
    if not masks:
        raise ValueError("输入的 masks 列表不能为空")

    # 把每张图都转成 0/1 numpy 数组
    bin_arrays = []
    for m in masks:
        arr = np.array(m.convert('L'), dtype=np.uint8)
        bin_arr = (arr > 127).astype(np.bool_)
        bin_arrays.append(bin_arr)

    # 做逐像素逻辑或
    union_bool = np.logical_or.reduce(bin_arrays)

    # 恢复成 0/255 uint8
    union_arr = union_bool.astype(np.uint8) * 255

    # 转回 PIL（二值）
    return Image.fromarray(union_arr, mode='L').convert('1')

def intersect_masks_np(masks: List[Image.Image]) -> Image.Image:
    """
    接受一个 PIL.Image 列表（mode='1' 或 'L' 的二值图），
    返回它们的交集。
    """
    if not masks:
        raise ValueError("输入的 masks 列表不能为空")

    # 把每张图都转成 0/1 numpy 数组
    bin_arrays = []
    for m in masks:
        arr = np.array(m.convert('L'), dtype=np.uint8)
        bin_arr = (arr > 127).astype(np.bool_)
        bin_arrays.append(bin_arr)

    # 做逐像素逻辑或
    intersect_bool = np.logical_and.reduce(bin_arrays)

    # 恢复成 0/255 uint8
    intersect_arr = intersect_bool.astype(np.uint8) * 255

    # 转回 PIL（二值）
    return Image.fromarray(intersect_arr, mode='L').convert('1')

def close_small_holes(pil_mask, kernel_size=5):
    """
    用闭运算填平小的黑点。
    kernel_size: 结构元尺寸，越大能填的洞越大，通常取奇数。
    """
    # 1. 转成 0/255 二值
    mask = np.array(pil_mask.convert('L'))
    _, bin_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    
    # 2. 定义结构元
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
    # 3. 闭运算
    closed = cv2.morphologyEx(bin_mask, cv2.MORPH_CLOSE, kernel)
    
    return Image.fromarray(closed)


def get_mask(src_image, tgt_image, threshold=1):
    """
    差异大的地方（差值大）当成前景（白），否则当背景（黑）
    """
    diff = ImageChops.difference(src_image, tgt_image)
    diff_gray = diff.convert("L")
    mask = diff_gray.point(lambda x: 255 if x >= threshold else 0).convert("1")
    return mask

def filter_small_components(pil_mask, area_threshold=0.10):
    """
    删除小于 area_threshold （默认 10%）的连通白色区域。
    pil_mask: PIL.Image，mode='L' 或 '1'（0/255 二值图）
    area_threshold: 阈值，相对于整张图面积的比例
    返回: 处理后的 PIL.Image
    """
    # 1. 转为二值 NumPy 数组（0,255）
    mask = np.array(pil_mask.convert('L'))
    # 确保是 0/255
    _, bin_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    
    # 2. 连通组件标记（4 或 8 邻域都可以）
    num_labels, labels = cv2.connectedComponents(bin_mask, connectivity=8)
    
    h, w = bin_mask.shape
    total_area = h * w
    total_area = np.count_nonzero(bin_mask)
    
    # 3. 遍历各个连通块
    output = np.zeros_like(bin_mask)
    for lbl in range(1, num_labels):  # 0 是背景
        # 取出该连通块
        comp_mask = (labels == lbl)
        comp_area = comp_mask.sum()
        # 面积比
        if comp_area >= area_threshold * total_area:
            # 保留
            output[comp_mask] = 255
    
    # 4. 转回 PIL
    return Image.fromarray(output)



def is_binary_255(t: torch.Tensor) -> bool:
    """
    判断给定的 tensor 是否只包含 0 和 255 两种值。
    """
    unique_vals = torch.unique(t)
    return torch.equal(unique_vals, torch.tensor([0], dtype=t.dtype)) or \
           torch.equal(unique_vals, torch.tensor([255], dtype=t.dtype)) or \
           torch.equal(unique_vals, torch.tensor([0, 255], dtype=t.dtype))

def get_weight(mask_u_ds, weight_type='log'):
    mask_u_ds_tensor = torch.from_numpy(np.array(mask_u_ds)).float()
    assert is_binary_255(mask_u_ds_tensor), "is_binary_255(mask_u_ds_tensor)"
    mask_u_ds_tensor_bool = mask_u_ds_tensor.bool()
    x = mask_u_ds_tensor_bool.numel() / mask_u_ds_tensor_bool.sum()
    if weight_type == 'log':
        weight = torch.log2(x) + 1
    elif weight_type == 'exp':
        weight = 2 ** (x**0.5 - 1)
    else:
        raise NotImplementedError(f'Support log | exp, but found {weight_type}')
    weight = torch.round(weight, decimals=6)
    assert weight >= 1, \
        f"weight >= 1 but {weight}, {mask_u_ds_tensor_bool.shape}, mask_u_ds_tensor_bool.numel(): {mask_u_ds_tensor_bool.numel()}, mask_u_ds_tensor_bool.sum(): {mask_u_ds_tensor_bool.sum()}"
    mask_u_ds_tensor[mask_u_ds_tensor==255] = weight
    mask_u_ds_tensor[mask_u_ds_tensor==0] = 1.0
    return mask_u_ds_tensor.unsqueeze(0)  # h w -> 1 h w

def get_weight_mask(pil_pixel_values, prompt=None, weight_type='log', need_weight='true'):
    # area_threshold = 1/64
    area_threshold = 0.001
    # base_kernel_size_factor = (5 / 448) ** 2
    # if len(pil_pixel_values) > 0:
    #     w, h = pil_pixel_values[-1].size
    #     kernel_size = max(int((base_kernel_size_factor * h * w) ** 0.5), 3)
    # else:
    kernel_size = 5

    if need_weight.lower() == 'false':
        mask_intersect = create_all_white_like(pil_pixel_values[-1])
        mask_intersect_ds = downsample_mask_pytorch(mask_intersect, factor=8)  # factor is downsample ratio of vae
        mask_intersect_ds = close_small_holes(mask_intersect_ds, kernel_size=kernel_size)
        weight = get_weight(mask_intersect_ds, weight_type)
        return mask_intersect_ds, weight

    filtered_masks = []
    for ii, j in enumerate(pil_pixel_values[:-1]):
        # each reference image will compare with target image to get mask
        mask = get_mask(j, pil_pixel_values[-1], threshold=18)
        # fill small holes
        fill_mask = close_small_holes(mask, kernel_size=kernel_size)
        # del small components
        filtered_mask = filter_small_components(fill_mask, area_threshold=0.3)
        # filtered_mask = fill_mask
        filtered_masks.append(filtered_mask)
    if len(filtered_masks) == 0:
        # t2i task do not have reference image
        assert len(pil_pixel_values) == 1, "len(pil_pixel_values) == 1"
        mask_intersect = create_all_white_like(pil_pixel_values[-1])
    else:
        mask_intersect = intersect_masks_np(filtered_masks)
    # while area / total area muse greater than 1/16 (just a threshold)
    mask_intersect_area_ratio = np.array(mask_intersect).astype(np.float32).sum() / np.prod(np.array(mask_intersect).shape)
    # print(mask_intersect_area_ratio)
    if mask_intersect_area_ratio < area_threshold:
        if mask_intersect_area_ratio == 0.0:
            # mask_intersect_area_ratio == 0 mean reconstruct data in stage 1
            assert len(pil_pixel_values) == 2, "len(pil_pixel_values) == 2"
            mask_intersect = create_all_white_like(pil_pixel_values[-1])
        else:
            # concat_images_row(pil_pixel_values + [mask_intersect], bg_color=(255,255,255)).show()
            raise ValueError(f'TOO SMALL mask_intersect_area_ratio: {mask_intersect_area_ratio}, prompt: {prompt}')
    mask_intersect_ds = downsample_mask_pytorch(mask_intersect, factor=8)  # factor is downsample ratio of vae
    mask_intersect_ds = close_small_holes(mask_intersect_ds, kernel_size=kernel_size)
    weight = get_weight(mask_intersect_ds, weight_type)
    return mask_intersect_ds, weight



def get_weight_mask_test(pil_pixel_values, prompt=None, weight_type='log'):
    area_threshold = 1/64
    base_kernel_size_factor = (5 / 448) ** 2
    if len(pil_pixel_values) > 0:
        w, h = pil_pixel_values[-1].size
        kernel_size = max(int((base_kernel_size_factor * h * w) ** 0.5), 3)
    else:
        kernel_size = 5

    filtered_masks = []
    for ii, j in enumerate(pil_pixel_values[:-1]):
        # each reference image will compare with target image to get mask
        mask = get_mask(j, pil_pixel_values[-1], threshold=18)
        # fill small holes
        fill_mask = close_small_holes(mask, kernel_size=kernel_size)
        # del small components
        filtered_mask = filter_small_components(fill_mask, area_threshold=1/64)
        # filtered_mask = fill_mask
        filtered_masks.append(filtered_mask)
    if len(filtered_masks) == 0:
        # t2i task do not have reference image
        assert len(pil_pixel_values) == 1, "len(pil_pixel_values) == 1"
        mask_intersect = create_all_white_like(pil_pixel_values[-1])
    else:
        mask_intersect = intersect_masks_np(filtered_masks)
    # while area / total area muse greater than 1/16 (just a threshold)
    mask_intersect_area_ratio = np.array(mask_intersect).astype(np.float32).sum() / np.prod(np.array(mask_intersect).shape)
    return mask_intersect_area_ratio