parse_utils/automasker.py

from PIL import Image
import numpy as np
import cv2
from typing import Union
from scipy.spatial import ConvexHull
from skimage.draw import polygon

DENSE_INDEX_MAP = {
    "background": [0],
    "torso": [1, 2],
    "right hand": [3],
    "left hand": [4],
    "right foot": [5],
    "left foot": [6],
    "right thigh": [7, 9],
    "left thigh": [8, 10],
    "right leg": [11, 13],
    "left leg": [12, 14],
    "left big arm": [15, 17],
    "right big arm": [16, 18],
    "left forearm": [19, 21],
    "right forearm": [20, 22],
    "face": [23, 24],
    "thighs": [7, 8, 9, 10],
    "legs": [11, 12, 13, 14],
    "hands": [3, 4],
    "feet": [5, 6],
    "big arms": [15, 16, 17, 18],
    "forearms": [19, 20, 21, 22],
}
ATR_MAPPING = {
    "Background": 0,
    "Hat": 1,
    "Hair": 2,
    "Sunglasses": 3,
    "Upper-clothes": 4,
    "Skirt": 5,
    "Pants": 6,
    "Dress": 7,
    "Belt": 8,
    "Left-shoe": 9,
    "Right-shoe": 10,
    "Face": 11,
    "Left-leg": 12,
    "Right-leg": 13,
    "Left-arm": 14,
    "Right-arm": 15,
    "Bag": 16,
    "Scarf": 17,
}
LIP_MAPPING = {
    "Background": 0,
    "Hat": 1,
    "Hair": 2,
    "Glove": 3,
    "Sunglasses": 4,
    "Upper-clothes": 5,
    "Dress": 6,
    "Coat": 7,
    "Socks": 8,
    "Pants": 9,
    "Jumpsuits": 10,
    "Scarf": 11,
    "Skirt": 12,
    "Face": 13,
    "Left-arm": 14,
    "Right-arm": 15,
    "Left-leg": 16,
    "Right-leg": 17,
    "Left-shoe": 18,
    "Right-shoe": 19,
}

PROTECT_BODY_PARTS = {
    "upper": ["Left-leg", "Right-leg"],
    "lower": ["Right-arm", "Left-arm", "Face"],
    "overall": [],
    "inner": ["Left-leg", "Right-leg"],
    "outer": ["Left-leg", "Right-leg"],
}
PROTECT_CLOTH_PARTS = {
    "upper": {"ATR": ["Skirt", "Pants"], "LIP": ["Skirt", "Pants"]},
    "lower": {
        "ATR": ["Upper-clothes", "Left-shoe", "Right-shoe"],
        "LIP": ["Upper-clothes", "Coat", "Left-shoe", "Right-shoe"],
    },
    "overall": {
        # 'ATR': [],
        # 'LIP': []
        "ATR": ["Left-shoe", "Right-shoe"],
        "LIP": ["Left-shoe", "Right-shoe"],
    },
    "inner": {
        "ATR": ["Dress", "Coat", "Skirt", "Pants"],
        "LIP": ["Dress", "Coat", "Skirt", "Pants", "Jumpsuits"],
    },
    "outer": {
        "ATR": ["Dress", "Pants", "Skirt"],
        "LIP": ["Upper-clothes", "Dress", "Pants", "Skirt", "Jumpsuits"],
    },
}
PUBLIC_ACCESSORY_PARTS = [
    "Hat",
    "Glove",
    "Sunglasses",
    "Scarf",
    "Bag",
    "Socks",
]  # 'Left-shoe', 'Right-shoe',
MASK_CLOTH_PARTS = {
    "upper": ["Upper-clothes", "Coat", "Dress", "Jumpsuits"],
    "lower": ["Pants", "Skirt", "Dress", "Jumpsuits"],
    "overall": [
        "Upper-clothes",
        "Dress",
        "Pants",
        "Skirt",
        "Coat",
        "Jumpsuits",
    ],  # , 'Left-shoe', 'Right-shoe'
    "inner": ["Upper-clothes"],
    "outer": [
        "Coat",
    ],
}
MASK_DENSE_PARTS = {
    "upper": ["torso", "big arms", "forearms"],
    "lower": ["thighs", "legs"],
    "overall": ["torso", "thighs", "legs", "big arms", "forearms"],
    "inner": ["torso"],
    "outer": ["torso", "big arms", "forearms"],
}


def random_convex_mask(height, width, min_area_ratio=0.25, jitter=10, extra_points=20):
    """
    在图像中随机位置生成一个凸多边形掩码，其包含一个面积不少于总面积1/4的扰动矩形。

    参数:
        height (int): 图像高度
        width (int): 图像宽度
        min_area_ratio (float): 最小面积比例
        jitter (int): 扰动范围（用于扩展凸多边形）
        extra_points (int): 随机扰动点数量

    返回:
        np.ndarray: 掩码数组 (0 背景, 1 前景)
    """
    rng = np.random.default_rng()
    total_area = height * width
    min_area = total_area * min_area_ratio

    # 随机生成一个矩形尺寸，满足面积要求
    aspect_ratio = rng.uniform(0.5, 2.0)  # 宽高比 [0.5, 2]
    rect_h = int(np.sqrt(min_area / aspect_ratio))
    rect_w = int(rect_h * aspect_ratio)

    # 随机偏移（使矩形尽量留在图像内）
    max_x0 = width - rect_w - 1
    max_y0 = height - rect_h - 1
    x0 = rng.integers(0, max(1, max_x0))
    y0 = rng.integers(0, max(1, max_y0))
    x1 = x0 + rect_w
    y1 = y0 + rect_h

    # 矩形四角 + 扰动点
    base_points = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]])
    jitter_points = rng.integers(
        low=[max(0, x0 - jitter), max(0, y0 - jitter)],
        high=[min(width, x1 + jitter), min(height, y1 + jitter)],
        size=(extra_points, 2),
    )
    all_points = np.vstack([base_points, jitter_points])

    # 计算凸包并生成掩码
    hull = ConvexHull(all_points)
    hull_points = all_points[hull.vertices]

    mask = np.zeros((height, width), dtype=np.uint8)
    rr, cc = polygon(hull_points[:, 1], hull_points[:, 0], shape=mask.shape)
    mask[rr, cc] = 1

    return mask


def part_mask_of(part: Union[str, list], parse: np.ndarray, mapping: dict):
    if isinstance(part, str):
        part = [part]
    mask = np.zeros_like(parse)
    for _ in part:
        if _ not in mapping:
            continue
        if isinstance(mapping[_], list):
            for i in mapping[_]:
                mask += parse == i
        else:
            mask += parse == mapping[_]
    return mask


def hull_mask(mask_area: np.ndarray):
    if len(mask_area.shape) > 2:
        mask_area = mask_area[:, :, 0]

    ret, binary = cv2.threshold(mask_area, 127, 255, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(
        binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    hull_mask = np.zeros_like(mask_area)
    for c in contours:
        hull = cv2.convexHull(c)
        hull_mask = cv2.fillPoly(np.zeros_like(mask_area), [hull], 255) | hull_mask
    return hull_mask


def create_square_mask(mask_area, expand_ratio=0.1):
    """
    将非零区域的最小外接矩形填充为1，并随机扩大边界

    Args:
        mask_area (np.ndarray): 原始 mask
        expand_ratio (float, optional): 边界随机扩大的最大比例。默认为0.1（10%）

    Returns:
        np.ndarray: 原始大小的 mask，矩形区域为1
    """
    # 找到非零区域
    rows = np.any(mask_area, axis=1)
    cols = np.any(mask_area, axis=0)

    # 获取非零区域边界
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]

    # 计算高度和宽度
    height = rmax - rmin + 1
    width = cmax - cmin + 1

    # 随机扩大边界
    expand_height = int(height * expand_ratio)
    expand_width = int(width * expand_ratio)

    # 计算新的边界，确保不超过图像尺寸
    new_rmin = max(0, rmin - np.random.randint(0, expand_height + 1))
    new_rmax = min(
        mask_area.shape[0] - 1, rmax + np.random.randint(0, expand_height + 1)
    )
    new_cmin = max(0, cmin - np.random.randint(0, expand_width + 1))
    new_cmax = min(
        mask_area.shape[1] - 1, cmax + np.random.randint(0, expand_width + 1)
    )

    # 创建新的 mask，并将矩形区域填充为1
    new_mask_area = np.zeros_like(mask_area)
    new_mask_area[new_rmin : new_rmax + 1, new_cmin : new_cmax + 1] = 1

    return new_mask_area


def create_bounding_box_mask(
    mask_area: np.ndarray,
    strong_protect_area: np.ndarray,
    dilate_kernel: np.ndarray,
    horizon_expand: bool = False,
) -> np.ndarray:
    """创建包含原始mask区域的外接矩形Mask，并排除保护区域

    Args:
        mask_area (np.ndarray): 原始mask区域
        strong_protect_area (np.ndarray): 需要排除的强保护区域
        dilate_kernel (np.ndarray): 用于膨胀操作的kernel
        horizon_expand (bool): 是否启用随机水平扩展

    Returns:
        np.ndarray: 处理后的矩形mask
    """
    # 找到mask区域的非零点坐标
    coords = cv2.findNonZero(mask_area)
    if coords is not None:
        # 获取外接矩形的坐标
        x, y, w, h = cv2.boundingRect(coords)

        # 随机水平扩展
        if horizon_expand:
            img_width = mask_area.shape[1]
            # 左右两边独立随机扩展 0.0～0.1倍的宽度
            left_expand = int(np.random.uniform(0.0, 0.2) * w)
            right_expand = int(np.random.uniform(0.0, 0.2) * w)

            # 计算扩展后的坐标，确保不超出图像边界
            x_expanded = max(0, x - left_expand)
            w_expanded = min(img_width - x_expanded, w + left_expand + right_expand)

            x, w = x_expanded, w_expanded

        # 创建外接矩形mask
        rect_mask = np.zeros_like(mask_area)
        rect_mask[y : y + h, x : x + w] = 1
        # 排除强保护区域
        rect_mask = rect_mask & (~strong_protect_area)
        # 最后的膨胀操作
        rect_mask = cv2.dilate(rect_mask, dilate_kernel, iterations=1)
        return rect_mask
    return mask_area


def cloth_agnostic_mask(
    densepose_mask: np.ndarray,
    schp_lip_mask: np.ndarray,
    schp_atr_mask: np.ndarray,
    part: str = "overall",
    square_cloth_mask: bool = False,
    **kwargs,
) -> Image.Image:
    if part == "full" or part == "dresses":
        part = "overall"
    assert part in ["upper", "lower", "overall", "inner", "outer"], (
        f"part should be one of ['upper', 'lower', 'overall', 'inner', 'outer'], but got {part}"
    )
    w, h = densepose_mask.shape[:2]

    dilate_kernel = max(w, h) // 500
    dilate_kernel = dilate_kernel if dilate_kernel % 2 == 1 else dilate_kernel + 1
    dilate_kernel = np.ones((dilate_kernel, dilate_kernel), np.uint8)

    kernal_size = max(w, h) // 50
    kernal_size = kernal_size if kernal_size % 2 == 1 else kernal_size + 1

    # Strong Protect Area (Hands, Face, Accessory, Feet)
    hands_protect_area = part_mask_of(
        ["hands", "feet"], densepose_mask, DENSE_INDEX_MAP
    )
    hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
    hands_protect_area = hands_protect_area & (
        part_mask_of(
            ["Left-arm", "Right-arm", "Left-leg", "Right-leg"],
            schp_atr_mask,
            ATR_MAPPING,
        )
        | part_mask_of(
            ["Left-arm", "Right-arm", "Left-leg", "Right-leg"],
            schp_lip_mask,
            LIP_MAPPING,
        )
    )
    face_protect_area = part_mask_of("Face", schp_lip_mask, LIP_MAPPING)

    strong_protect_area = hands_protect_area | face_protect_area

    # Weak Protect Area (Hair, Irrelevant Clothes, Body Parts)
    body_protect_area = part_mask_of(
        PROTECT_BODY_PARTS[part], schp_lip_mask, LIP_MAPPING
    ) | part_mask_of(PROTECT_BODY_PARTS[part], schp_atr_mask, ATR_MAPPING)
    hair_protect_area = part_mask_of(
        ["Hair"], schp_lip_mask, LIP_MAPPING
    ) | part_mask_of(["Hair"], schp_atr_mask, ATR_MAPPING)
    cloth_protect_area = part_mask_of(
        PROTECT_CLOTH_PARTS[part]["LIP"], schp_lip_mask, LIP_MAPPING
    ) | part_mask_of(PROTECT_CLOTH_PARTS[part]["ATR"], schp_atr_mask, ATR_MAPPING)
    accessory_protect_area = part_mask_of(
        PUBLIC_ACCESSORY_PARTS, schp_lip_mask, LIP_MAPPING
    ) | part_mask_of(PUBLIC_ACCESSORY_PARTS, schp_atr_mask, ATR_MAPPING)
    weak_protect_area = (
        body_protect_area
        | cloth_protect_area
        | hair_protect_area
        | strong_protect_area
        | accessory_protect_area
    )

    # Mask Area
    strong_mask_area = part_mask_of(
        MASK_CLOTH_PARTS[part], schp_lip_mask, LIP_MAPPING
    ) | part_mask_of(MASK_CLOTH_PARTS[part], schp_atr_mask, ATR_MAPPING)
    strong_mask_area = cv2.dilate(
        strong_mask_area, dilate_kernel // 2 + 1, iterations=1
    )  # ADD: max pooling

    background_area = part_mask_of(
        ["Background"], schp_lip_mask, LIP_MAPPING
    ) & part_mask_of(["Background"], schp_atr_mask, ATR_MAPPING)
    mask_dense_area = part_mask_of(
        MASK_DENSE_PARTS[part] + ['right foot', 'left foot'], densepose_mask, DENSE_INDEX_MAP
    )
    mask_dense_area = cv2.resize(
        mask_dense_area.astype(np.uint8),
        None,
        fx=0.25,
        fy=0.25,
        interpolation=cv2.INTER_NEAREST,
    )
    mask_dense_area = cv2.dilate(mask_dense_area, dilate_kernel, iterations=2)
    mask_dense_area = cv2.resize(
        mask_dense_area.astype(np.uint8),
        None,
        fx=4,
        fy=4,
        interpolation=cv2.INTER_NEAREST,
    )

    mask_area = (
        np.ones_like(densepose_mask) & (~weak_protect_area) & (~background_area)
    ) | mask_dense_area

    # 边缘平滑（去除毛刺）
    mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
    mask_area[mask_area < 100] = 0
    mask_area[mask_area >= 100] = 1

    # 确保只有一个连通域
    num_labels, labels = cv2.connectedComponents(mask_area.astype(np.uint8))
    if num_labels > 2:  # 背景(0)和一个前景区域
        label_counts = np.bincount(labels.flatten())  # 找到最大连通区域的标签
        label_counts[0] = 0  # 排除背景
        largest_label = np.argmax(label_counts)
        mask_area = (labels == largest_label).astype(
            np.uint8
        )  # 创建只包含最大连通区域的掩码，并填充内部空洞

    # 凸包扩张
    mask_area = hull_mask(mask_area * 255) // 255  # Convex Hull to expand the mask area
    mask_area = mask_area & (~weak_protect_area)
    mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
    mask_area[mask_area < 25] = 0
    mask_area[mask_area >= 25] = 1

    # 创建方形mask
    if square_cloth_mask:  # v2
        weak_protect_area_ = weak_protect_area & (~mask_area)  # 防止边缘过度保护
        mask_area = create_bounding_box_mask(
            mask_area, weak_protect_area_, dilate_kernel
        )
        mask_area = mask_area & ~face_protect_area

    # Pooling
    mask_area = cv2.dilate(mask_area, dilate_kernel - 2, iterations=2)

    return Image.fromarray(mask_area * 255)


def multi_ref_cloth_agnostic_mask(
    densepose_mask: np.ndarray,
    schp_lip_mask: np.ndarray,
    schp_atr_mask: np.ndarray,
    square_cloth_mask: bool = False,
    horizon_expand: bool = False,
    **kwargs,
) -> Image.Image:
    w, h = densepose_mask.shape[:2]

    dilate_kernel = max(w, h) // 500
    dilate_kernel = dilate_kernel if dilate_kernel % 2 == 1 else dilate_kernel + 1
    dilate_kernel = np.ones((dilate_kernel, dilate_kernel), np.uint8)

    kernal_size = max(w, h) // 50
    kernal_size = kernal_size if kernal_size % 2 == 1 else kernal_size + 1

    # Strong Protect Area (Hands, Face, Accessory, Feet)
    # hands_protect_area = part_mask_of(['hands', 'feet'], densepose_mask, DENSE_INDEX_MAP)
    # hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1)
    # hands_protect_area = hands_protect_area & \
    #     (part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_atr_mask, ATR_MAPPING) | \
    #         part_mask_of(['Left-arm', 'Right-arm', 'Left-leg', 'Right-leg'], schp_lip_mask, LIP_MAPPING))
    face_protect_area = part_mask_of("Face", schp_lip_mask, LIP_MAPPING)
    # strong_protect_area = hands_protect_area | face_protect_area
    strong_protect_area = face_protect_area

    # Mask Area
    mask_keys = [
        "Upper-clothes",
        "Dress",
        "Pants",
        "Skirt",
        "Coat",
        "Jumpsuits",
        "Left-shoe",
        "Right-shoe",
        "Bag",
        "Socks",
        "Belt",
    ]
    strong_mask_area = part_mask_of(
        mask_keys, schp_lip_mask, LIP_MAPPING
    ) | part_mask_of(mask_keys, schp_atr_mask, ATR_MAPPING)

    strong_mask_area = cv2.dilate(
        strong_mask_area, dilate_kernel // 2 + 1, iterations=1
    )  # ADD: max pooling
    mask_dense_area = part_mask_of(
        MASK_DENSE_PARTS["overall"] + ['right foot', 'left foot'], densepose_mask, DENSE_INDEX_MAP
    )
    mask_dense_area = cv2.resize(
        mask_dense_area.astype(np.uint8),
        None,
        fx=0.25,
        fy=0.25,
        interpolation=cv2.INTER_NEAREST,
    )
    mask_dense_area = cv2.dilate(mask_dense_area, dilate_kernel, iterations=2)
    mask_dense_area = cv2.resize(
        mask_dense_area.astype(np.uint8),
        None,
        fx=4,
        fy=4,
        interpolation=cv2.INTER_NEAREST,
    )

    mask_area = (strong_mask_area | mask_dense_area) & (~strong_protect_area)

    mask_area = cv2.dilate(
        mask_area, dilate_kernel // 2 + 1, iterations=1
    )  # ADD: max pooling

    # 边缘平滑（去除毛刺）
    mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
    mask_area[mask_area < 100] = 0
    mask_area[mask_area >= 100] = 1

    # 确保只有一个连通域
    # num_labels, labels = cv2.connectedComponents(mask_area.astype(np.uint8))
    # if num_labels > 2:  # 背景(0)和一个前景区域
    #     label_counts = np.bincount(labels.flatten())  # 找到最大连通区域的标签
    #     label_counts[0] = 0  # 排除背景
    #     largest_label = np.argmax(label_counts)
    #     mask_area = (labels == largest_label).astype(
    #         np.uint8
    #     )  # 创建只包含最大连通区域的掩码，并填充内部空洞

    # 凸包扩张
    mask_area = hull_mask(mask_area * 255) // 255  # Convex Hull to expand the mask area
    mask_area = cv2.GaussianBlur(mask_area * 255, (kernal_size, kernal_size), 0)
    mask_area[mask_area < 25] = 0
    mask_area[mask_area >= 25] = 1

    # 创建方形mask
    if square_cloth_mask:  # v2
        # weak_protect_area_ = weak_protect_area & (~mask_area) # 防止边缘过度保护
        mask_area = create_bounding_box_mask(
            mask_area, strong_protect_area, dilate_kernel, horizon_expand
        )
        mask_area = mask_area & ~strong_protect_area

    # Pooling
    mask_area = cv2.dilate(mask_area, dilate_kernel - 2, iterations=2)

    return Image.fromarray(mask_area * 255)