| import numpy as np |
| import cv2 |
| from PIL import Image |
| from typing import Union |
| from mappings import ( |
| DENSE_INDEX_MAP, |
| ATR_MAPPING, |
| LIP_MAPPING, |
| MASK_CLOTH_PARTS, |
| MASK_DENSE_PARTS, |
| PROTECT_BODY_PARTS, |
| PROTECT_CLOTH_PARTS |
| ) |
|
|
| def get_person_height(mask): |
| mask = (mask > 127).astype(np.uint8) |
| nonzero_pixels = np.any(mask, axis=1) |
| if not np.any(nonzero_pixels): |
| return 50 |
| ymin, ymax = np.where(nonzero_pixels)[0][[0, -1]] |
| person_height = max(50, ymax - ymin) |
| return person_height |
|
|
| def calculate_kernels(person_height): |
| dilate_size = max(1, person_height // 30) |
| dilate_size = dilate_size if dilate_size % 2 == 1 else dilate_size + 1 |
| dilate_kernel = np.ones((dilate_size, dilate_size), np.uint8) |
| kernel_size = max(5, person_height // 15) |
| kernel_size = kernel_size if kernel_size % 2 == 1 else kernel_size + 1 |
| return dilate_kernel, kernel_size |
|
|
| def part_mask_of(part: Union[str, list], parse: np.ndarray, mapping: dict): |
| if isinstance(part, str): |
| part = [part] |
| mask = np.zeros_like(parse, dtype=np.uint8) |
| for p in part: |
| if p not in mapping: |
| continue |
| val = mapping[p] |
| if isinstance(val, list): |
| for i in val: |
| mask[parse == i] = 1 |
| else: |
| mask[parse == val] = 1 |
| return mask |
|
|
| def hull_mask(mask_area: np.ndarray): |
| mask_area = (mask_area * 255).astype(np.uint8) |
| _, binary = cv2.threshold(mask_area, 127, 255, cv2.THRESH_BINARY) |
| contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| hull_mask_result = np.zeros_like(binary) |
| for cnt in contours: |
| hull = cv2.convexHull(cnt) |
| cv2.fillPoly(hull_mask_result, [hull], 255) |
| return (hull_mask_result > 127).astype(np.uint8) |
|
|
| def compute_strong_protect_area(densepose_mask, schp_lip_mask, schp_atr_mask, dilate_kernel): |
| hands_protect_area = part_mask_of(['hands', 'feet', 'face'], densepose_mask, DENSE_INDEX_MAP) |
| hands_protect_area = cv2.dilate(hands_protect_area, dilate_kernel, iterations=1) |
| arms_legs_schp = ( |
| 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) |
| ) |
| hands_protect_area = hands_protect_area & arms_legs_schp |
| face_protect_area = part_mask_of('Face', schp_lip_mask, LIP_MAPPING) |
| return (hands_protect_area | face_protect_area).astype(np.uint8) |
|
|
| def compute_weak_protect_area(schp_lip_mask, schp_atr_mask, strong_protect_area, part): |
| 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]['ATR'], schp_atr_mask, ATR_MAPPING) |
| accessory_parts = ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks'] |
| accessory_protect_area = ( |
| part_mask_of(accessory_parts, schp_lip_mask, LIP_MAPPING) | |
| part_mask_of(accessory_parts, schp_atr_mask, ATR_MAPPING) |
| ) |
| weak_area = (body_protect_area | cloth_protect_area | hair_protect_area | |
| strong_protect_area | accessory_protect_area) |
| return weak_area.astype(np.uint8) |
|
|
| def compute_mask_area(densepose_mask, schp_lip_mask, schp_atr_mask, |
| weak_protect_area, strong_protect_area, |
| dilate_kernel, part): |
| 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) |
| ) |
| 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], 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) |
| target_shape = densepose_mask.shape |
| def resize_mask(mask, target_shape): |
| return cv2.resize(mask.astype(np.uint8), (target_shape[1], target_shape[0]), interpolation=cv2.INTER_NEAREST) |
| weak_protect_area = resize_mask(weak_protect_area, target_shape) |
| background_area = resize_mask(background_area, target_shape) |
| mask_dense_area = resize_mask(mask_dense_area, target_shape) |
| base_area = np.ones_like(densepose_mask, dtype=np.uint8) |
| mask_area = (base_area & (~weak_protect_area) & (~background_area)) | mask_dense_area |
| mask_area = hull_mask(mask_area) |
| mask_area = (mask_area & (~weak_protect_area)).astype(np.uint8) |
| return mask_area, background_area |
|
|
| def safe_dilate(mask, kernel, iterations, protect_mask): |
| dilated_mask = mask.copy() |
| for _ in range(iterations): |
| temp_dilated = cv2.dilate(dilated_mask, kernel, iterations=1) |
| temp_dilated[protect_mask > 0] = 0 |
| dilated_mask = temp_dilated |
| return dilated_mask |
|
|
| def finalize_mask(mask_area, kernal_size, strong_mask_area, strong_protect_area, dilate_kernel): |
| mask_area = (mask_area * 255).astype(np.uint8) |
| mask_area = cv2.GaussianBlur(mask_area, (kernal_size, kernal_size), 0) |
| mask_area = np.where(mask_area < 25, 0, 1).astype(np.uint8) |
| mask_area = (mask_area | strong_mask_area) & (~strong_protect_area) |
| mask_area = safe_dilate(mask_area, dilate_kernel, iterations=1, protect_mask=strong_protect_area) |
| return Image.fromarray((mask_area * 255).astype(np.uint8)) |
|
|
| def keep_largest_connected_component(mask_np): |
| mask_uint8 = (mask_np > 128).astype(np.uint8) |
| num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_uint8, connectivity=8) |
| if num_labels <= 1: |
| return mask_uint8 * 255 |
| largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA]) |
| cleaned_mask = np.zeros_like(mask_uint8) |
| cleaned_mask[labels == largest_label] = 1 |
| return cleaned_mask * 255 |
