import os from PIL import Image from typing import Union import numpy as np import cv2 from diffusers.image_processor import VaeImageProcessor import torch from model.SCHP import SCHP # type: ignore from model.DensePose import DensePose # type: ignore 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'], 'LIP': ['Upper-clothes', 'Coat'] }, 'overall': { 'ATR': [], 'LIP': [] }, 'inner': { 'ATR': ['Dress', 'Coat', 'Skirt', 'Pants'], 'LIP': ['Dress', 'Coat', 'Skirt', 'Pants', 'Jumpsuits'] }, 'outer': { 'ATR': ['Dress', 'Pants', 'Skirt'], 'LIP': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Jumpsuits'] } } MASK_CLOTH_PARTS = { 'upper': ['Upper-clothes', 'Coat', 'Dress', 'Jumpsuits'], 'lower': ['Pants', 'Skirt', 'Dress', 'Jumpsuits'], 'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits'], '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'] } schp_public_protect_parts = ['Hat', 'Hair', 'Sunglasses', 'Left-shoe', 'Right-shoe', 'Bag', 'Glove', 'Scarf'] schp_protect_parts = { 'upper': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits'], 'lower': ['Left-arm', 'Right-arm', 'Upper-clothes', 'Coat'], 'overall': [], 'inner': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits', 'Coat'], 'outer': ['Left-leg', 'Right-leg', 'Skirt', 'Pants', 'Jumpsuits', 'Upper-clothes'] } schp_mask_parts = { 'upper': ['Upper-clothes', 'Dress', 'Coat', 'Jumpsuits'], 'lower': ['Pants', 'Skirt', 'Dress', 'Jumpsuits', 'socks'], 'overall': ['Upper-clothes', 'Dress', 'Pants', 'Skirt', 'Coat', 'Jumpsuits', 'socks'], 'inner': ['Upper-clothes'], 'outer': ['Coat',] } dense_mask_parts = { 'upper': ['torso', 'big arms', 'forearms'], 'lower': ['thighs', 'legs'], 'overall': ['torso', 'thighs', 'legs', 'big arms', 'forearms'], 'inner': ['torso'], 'outer': ['torso', 'big arms', 'forearms'] } # 추가한 함수 def save_mask(mask: np.ndarray, file_path: str): # numpy 배열을 PIL 이미지로 변환 mask_image = Image.fromarray(mask.astype(np.uint8) * 255) # 이진 마스크인 경우 0과 1을 0과 255로 변환 mask_image.save(file_path) def vis_mask(image, mask): image = np.array(image).astype(np.uint8) mask = np.array(mask).astype(np.uint8) mask[mask > 127] = 255 mask[mask <= 127] = 0 mask = np.expand_dims(mask, axis=-1) mask = np.repeat(mask, 3, axis=-1) mask = mask / 255 return Image.fromarray((image * (1 - mask)).astype(np.uint8)) 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): 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 morph_open(mask: np.ndarray): kernel = np.ones((10, 10), np.uint8) #원래 25 25 but 커지면 좋지 않음. # Open operation 적용 cleaned_mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) # numpy 배열을 PIL 이미지로 변환 return cleaned_mask def morph_close(mask: np.ndarray): kernel = np.ones((80, 80), np.uint8) # CLOSE operation 적용 cleaned_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # numpy 배열을 PIL 이미지로 변환 return cleaned_mask def remove_bottom_part(mask: np.ndarray, y_threshold: int): """ 이미지의 y_threshold 아래에 있는 부분을 삭제. :param mask: 입력 마스크 (numpy 배열) :param y_threshold: 제거할 Y 좌표 값 :return: 수정된 마스크 (numpy 배열) """ # y_threshold 아래의 모든 픽셀을 0으로 설정 mask[y_threshold:, :] = 0 return mask class AutoMasker: def __init__( self, densepose_ckpt='./Models/DensePose', schp_ckpt='./Models/SCHP', device='cuda'): np.random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed(0) self.densepose_processor = DensePose(densepose_ckpt, device) self.schp_processor_atr = SCHP(ckpt_path=os.path.join(schp_ckpt, 'exp-schp-201908301523-atr.pth'), device=device) self.schp_processor_lip = SCHP(ckpt_path=os.path.join(schp_ckpt, 'exp-schp-201908261155-lip.pth'), device=device) self.mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True) def process_densepose(self, image_or_path): return self.densepose_processor(image_or_path, resize=1024) def process_schp_lip(self, image_or_path): return self.schp_processor_lip(image_or_path) def process_schp_atr(self, image_or_path): return self.schp_processor_atr(image_or_path) def preprocess_image(self, image_or_path): return { 'densepose': self.densepose_processor(image_or_path, resize=1024), 'schp_atr': self.schp_processor_atr(image_or_path), 'schp_lip': self.schp_processor_lip(image_or_path) } @staticmethod def cloth_agnostic_mask( densepose_mask: Image.Image, schp_lip_mask: Image.Image, schp_atr_mask: Image.Image, part: str='long dress', **kwargs ): assert part in ['long sleeve', 'short sleeve', 'long pants', 'short pants', 'long dress', 'short dress'], f"part should be one of ['long sleeve', 'short sleeve', 'long pants', 'short pants', 'long dress', 'short dress'], but got {part}" w, h = densepose_mask.size dilate_kernel = max(w, h) // 250 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) // 25 kernal_size = kernal_size if kernal_size % 2 == 1 else kernal_size + 1 densepose_mask = np.array(densepose_mask) schp_lip_mask = np.array(schp_lip_mask) schp_atr_mask = np.array(schp_atr_mask) # part_mask_of(MASK_CLOTH_PARTS[part], schp_atr_mask, ATR_MAPPING) # Mask Area if part == 'long sleeve': mask_area = part_mask_of(['torso', 'big arms', 'forearms'], densepose_mask, DENSE_INDEX_MAP) #완전 기본 마스크 # 튀어나온 부분 밀어버리기 height = mask_area.shape[0] y_threshold = int(height * 0.5) # 이미지 높이의 50퍼센트 이하. 50퍼센트가 딱 적당함. # 밑부분 제거 mask_area = remove_bottom_part(mask_area, y_threshold) # 모폴로지 연산 수행 dilate_kernel= np.ones((15, 15), np.uint8) mask_area = cv2.dilate(mask_area, dilate_kernel, iterations=1) mask_area = morph_open(mask_area) if part == 'short sleeve': mask_area = part_mask_of(['torso', 'big arms'], densepose_mask, DENSE_INDEX_MAP) #완전 기본 마스크 # 튀어나온 부분 밀어버리기 height = mask_area.shape[0] y_threshold = int(height * 0.5) # 이미지 높이의 50퍼센트 이하. 50퍼센트가 딱 적당함. # 밑부분 제거 mask_area = remove_bottom_part(mask_area, y_threshold) # 모폴로지 연산 수행 dilate_kernel= np.ones((15, 15), np.uint8) mask_area = cv2.dilate(mask_area, dilate_kernel, iterations=1) mask_area = morph_open(mask_area) # 바로 아래 코드에서 사용자가 이미지가 치마를 입고 있으면 atr_mask에 pants요소가 없어서 안 됨. # 즉 동적으로 이미지 입력이 있다면, -> 그 상황에서는 무조건, "long_dress에서 사용하는 default mask"를 사용하도록 만들어야 할 것임. if part == 'long pants': mask_area = (part_mask_of(['thighs', 'legs'], densepose_mask, DENSE_INDEX_MAP) | \ part_mask_of(['Pants'], schp_atr_mask, ATR_MAPPING)) if part == 'short pants': mask_area = (part_mask_of(['thighs'], densepose_mask, DENSE_INDEX_MAP)| \ part_mask_of(['Pants'], schp_atr_mask, ATR_MAPPING)) #바지 밑단 자르기 height = mask_area.shape[0] y_threshold = int(height * 0.67) # 이미지 높이의 50퍼센트 이하. 50퍼센트가 딱 적당함. # 밑부분 제거 mask_area = remove_bottom_part(mask_area, y_threshold) if part == 'long dress': # 기존 코드의 part == 'overall' 일때 해당하는 코드임. # 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([], schp_lip_mask, LIP_MAPPING) | part_mask_of([], 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['overall']['LIP'], schp_lip_mask, LIP_MAPPING) | \ part_mask_of(PROTECT_CLOTH_PARTS['overall']['ATR'], schp_atr_mask, ATR_MAPPING) accessory_protect_area = part_mask_of((accessory_parts := ['Hat', 'Glove', 'Sunglasses', 'Bag', 'Left-shoe', 'Right-shoe', 'Scarf', 'Socks']), schp_lip_mask, LIP_MAPPING) | \ part_mask_of(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['overall'], schp_lip_mask, LIP_MAPPING) | \ part_mask_of(MASK_CLOTH_PARTS['overall'], 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['overall'], 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 = 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_area = (mask_area | strong_mask_area) & (~strong_protect_area) mask_area = cv2.dilate(mask_area, dilate_kernel, iterations=1) #return Image.fromarray(mask_area * 255) if part == 'short dress': mask_area_sp = (part_mask_of(['thighs'], densepose_mask, DENSE_INDEX_MAP)| \ part_mask_of(['Pants'], schp_atr_mask, ATR_MAPPING)) #바지 밑단 자르기 height = mask_area_sp.shape[0] y_threshold = int(height * 0.67) # 이미지 높이의 67퍼 이하가 이 경우에 적당. # 밑부분 제거 mask_area_sp = remove_bottom_part(mask_area_sp, y_threshold) # 바지 갈라짐 없어지도록 모폴로지 연산 수행 mask_area_sp = morph_close(mask_area_sp) # closing이 메우는 핵심임. dilate_kernel= np.ones((15, 15), np.uint8) mask_area_sp = cv2.dilate(mask_area_sp, dilate_kernel, iterations=1) # 윗부분 mask_area_up = part_mask_of(['torso', 'big arms', 'forearms'], densepose_mask, DENSE_INDEX_MAP) mask_area_up = cv2.dilate(mask_area_up, dilate_kernel, iterations=1) # mask_area_up = morph_open(mask_area_up) #합치기 mask_area = mask_area_sp | mask_area_up return Image.fromarray(mask_area.astype(np.uint8) * 255) def __call__( self, image: Union[str, Image.Image], mask_type: str = "long sleeve", ): assert mask_type in ['long sleeve', 'short sleeve', 'long pants', 'short pants', 'long dress', 'short dress'], f"mask_type should be one of ['long sleeve', 'short sleeve', 'long pants', 'short pants', 'long dress', 'short dress], but got {mask_type}" preprocess_results = self.preprocess_image(image) mask = self.cloth_agnostic_mask( preprocess_results['densepose'], preprocess_results['schp_lip'], preprocess_results['schp_atr'], part=mask_type, ) return { 'mask': mask, 'densepose': preprocess_results['densepose'], 'schp_lip': preprocess_results['schp_lip'], 'schp_atr': preprocess_results['schp_atr'] } if __name__ == '__main__': pass