File size: 6,378 Bytes
436df5c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | 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 |