#!/usr/bin/env python # -*- encoding: utf-8 -*- """ @Author : Peike Li @Contact : peike.li@yahoo.com @File : simple_extractor.py @Time : 8/30/19 8:59 PM @Desc : Simple Extractor @License : This source code is licensed under the license found in the LICENSE file in the root directory of this source tree. """ import os import cv2 import torch import argparse import numpy as np from PIL import Image from tqdm import tqdm from torch.utils.data import DataLoader import torchvision.transforms as transforms import networks from utils.transforms import transform_logits from datasets.simple_extractor_dataset import SimpleFolderDataset CROP_PADDING = 12 dataset_settings = { 'lip': { 'input_size': [473, 473], 'num_classes': 20, 'label': ['Background', 'Hat', 'Hair', 'Glove', 'Sunglasses', 'Upper-clothes', 'Dress', 'Coat', 'Socks', 'Pants', 'Jumpsuits', 'Scarf', 'Skirt', 'Face', 'Left-arm', 'Right-arm', 'Left-leg', 'Right-leg', 'Left-shoe', 'Right-shoe'] }, 'atr': { 'input_size': [512, 512], 'num_classes': 18, 'label': ['Background', 'Hat', 'Hair', 'Sunglasses', 'Upper-clothes', 'Skirt', 'Pants', 'Dress', 'Belt', 'Left-shoe', 'Right-shoe', 'Face', 'Left-leg', 'Right-leg', 'Left-arm', 'Right-arm', 'Bag', 'Scarf'] }, 'pascal': { 'input_size': [512, 512], 'num_classes': 7, 'label': ['Background', 'Head', 'Torso', 'Upper Arms', 'Lower Arms', 'Upper Legs', 'Lower Legs'], } } def get_arguments(): """Parse all the arguments provided from the CLI. Returns: A list of parsed arguments. """ parser = argparse.ArgumentParser(description="Self Correction for Human Parsing") parser.add_argument("--dataset", type=str, default='lip', choices=['lip', 'atr', 'pascal']) parser.add_argument("--model-restore", type=str, default='', help="restore pretrained model parameters.") parser.add_argument("--gpu", type=str, default='0', help="choose gpu device.") parser.add_argument("--input-dir", type=str, default='', help="path of input image folder.") parser.add_argument("--output-dir", type=str, default='', help="path of output image folder.") parser.add_argument("--logits", action='store_true', default=False, help="whether to save the logits.") return parser.parse_args() def get_palette(num_cls): """ Returns the color map for visualizing the segmentation mask. Args: num_cls: Number of classes Returns: The color map """ n = num_cls palette = [0] * (n * 3) for j in range(0, n): lab = j palette[j * 3 + 0] = 0 palette[j * 3 + 1] = 0 palette[j * 3 + 2] = 0 i = 0 while lab: palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i)) palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i)) palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i)) i += 1 lab >>= 3 return palette def create_mask(parsing_result, class_ids): """특정 클래스만 흰색(mask)으로 만드는 함수""" mask = np.zeros_like(parsing_result).astype(np.uint8) for cid in class_ids: mask[parsing_result == cid] = 255 return mask def crop_to_alpha(image, alpha): """투명하지 않은 영역만 남기도록 이미지 여백 제거""" ys, xs = np.where(alpha > 0) if len(xs) == 0 or len(ys) == 0: return image height, width = alpha.shape x1 = max(xs.min() - CROP_PADDING, 0) y1 = max(ys.min() - CROP_PADDING, 0) x2 = min(xs.max() + CROP_PADDING + 1, width) y2 = min(ys.max() + CROP_PADDING + 1, height) return image[y1:y2, x1:x2] def extract_part(original_image, mask): """마스크를 적용하여 의류 부분 추출""" mask_3ch = cv2.merge([mask, mask, mask]) extracted = np.where( mask_3ch == 255, original_image, 0 ) rgba = cv2.cvtColor(extracted, cv2.COLOR_RGB2RGBA) rgba[:, :, 3] = mask return crop_to_alpha(rgba, mask) def main(): args = get_arguments() if args.gpu == 'None': raise RuntimeError('GPU 전용 실행만 허용됩니다. --gpu 값은 None이 될 수 없습니다.') gpus = [int(i) for i in args.gpu.split(',')] assert len(gpus) == 1 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu if not torch.cuda.is_available(): raise RuntimeError('CUDA GPU를 사용할 수 없습니다. GPU 전용 실행만 허용됩니다.') device = torch.device('cuda:0') num_classes = dataset_settings[args.dataset]['num_classes'] input_size = dataset_settings[args.dataset]['input_size'] label = dataset_settings[args.dataset]['label'] print("Evaluating total class number {} with {}".format(num_classes, label)) print("Using device:", device) model = networks.init_model('resnet101', num_classes=num_classes, pretrained=None) state_dict = torch.load(args.model_restore, map_location=device)['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v model.load_state_dict(new_state_dict) model.to(device) model.eval() transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229]) ]) dataset = SimpleFolderDataset(root=args.input_dir, input_size=input_size, transform=transform) dataloader = DataLoader(dataset) # Output 디렉토리 삭제 후 생성 if os.path.exists(args.output_dir): import shutil shutil.rmtree(args.output_dir) os.makedirs(args.output_dir) palette = get_palette(num_classes) # 의류별 클래스 정의 (LIP 데이터셋 기준) - 옷만 분류 clothing_classes = { 'hat': [1], # Hat 'upper_clothes': [5], # Upper-clothes 'dress': [6], # Dress 'coat': [7], # Coat 'pants': [9], # Pants 'jumpsuits': [10], # Jumpsuits 'scarf': [11], # Scarf 'skirt': [12], # Skirt } with torch.no_grad(): for idx, batch in enumerate(tqdm(dataloader)): image, meta = batch img_name = meta['name'][0] c = meta['center'].numpy()[0] s = meta['scale'].numpy()[0] w = meta['width'].numpy()[0] h = meta['height'].numpy()[0] output = model(image.to(device)) upsample = torch.nn.Upsample(size=input_size, mode='bilinear', align_corners=True) upsample_output = upsample(output[0][-1][0].unsqueeze(0)) upsample_output = upsample_output.squeeze() upsample_output = upsample_output.permute(1, 2, 0) # CHW -> HWC logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=input_size) parsing_result = np.argmax(logits_result, axis=2) if args.logits: logits_result_path = os.path.join(args.output_dir, img_name[:-4] + '.npy') np.save(logits_result_path, logits_result) # 의류별 추출 및 저장 print(f"의류별 추출 중: {img_name}") # 원본 이미지 로드 original_image_path = os.path.join(args.input_dir, img_name) original_image = cv2.imread(original_image_path) original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB) # 원본 이미지와 파싱 결과의 크기가 같도록 리사이징 original_image_resized = cv2.resize(original_image, (parsing_result.shape[1], parsing_result.shape[0])) # 의류별 확률 출력 print(f"\n=== {img_name} 의류 확률 ===") # 의류 부위별 추출 및 저장 for clothing_name, class_ids in clothing_classes.items(): clothing_mask = create_mask(parsing_result, class_ids) # 마스크가 비어있지 않은 경우만 처리 if clothing_mask.sum() > 0: # 점수 계산: 해당 의류의 평균 신뢰도 점수 class_id = class_ids[0] class_scores = logits_result[:, :, class_id] average_score = class_scores[clothing_mask > 0].mean() # Softmax 확률로 변환 (0-1 범위) softmax_prob = 1 / (1 + np.exp(-average_score)) print(f"{clothing_name:15} : {softmax_prob*100:6.2f}%") clothing_img = extract_part(original_image_resized, clothing_mask) clothing_output_path = os.path.join(args.output_dir, img_name[:-4] + f'_{clothing_name}.png') cv2.imwrite(clothing_output_path, cv2.cvtColor(clothing_img, cv2.COLOR_RGBA2BGRA)) else: print(f"{clothing_name:15} : 감지 안됨") return if __name__ == '__main__': main()