| |
| |
| |
| |
| |
| """ |
| @Author : Peike Li |
| @Contact : peike.li@yahoo.com |
| @File : simple_extractor.py |
| @Time : 8/30/19 8:59 PM |
| @Desc : Simple Extractor (modified for single image input) |
| """ |
|
|
| import os |
| import torch |
| import argparse |
| import numpy as np |
| from PIL import Image |
| from tqdm import tqdm |
| import cv2 |
|
|
| from torch.utils.data import Dataset, DataLoader |
| import torchvision.transforms as transforms |
|
|
| import networks |
| from preprocess.utils.transforms import transform_logits, get_affine_transform |
|
|
|
|
| class SimpleFileDataset(Dataset): |
| def __init__(self, image_path, input_size=[512, 512], transform=None): |
| self.image_path = image_path |
| self.input_size = np.asarray(input_size) |
| self.transform = transform |
| self.aspect_ratio = input_size[1] * 1.0 / input_size[0] |
| self.img_name = os.path.basename(image_path) |
|
|
| def __len__(self): |
| return 1 |
|
|
| def _box2cs(self, box): |
| x, y, w, h = box[:4] |
| return self._xywh2cs(x, y, w, h) |
|
|
| def _xywh2cs(self, x, y, w, h): |
| center = np.zeros((2), dtype=np.float32) |
| center[0] = x + w * 0.5 |
| center[1] = y + h * 0.5 |
| if w > self.aspect_ratio * h: |
| h = w * 1.0 / self.aspect_ratio |
| elif w < self.aspect_ratio * h: |
| w = h * self.aspect_ratio |
| scale = np.array([w, h], dtype=np.float32) |
| return center, scale |
|
|
| def __getitem__(self, index): |
| img = cv2.imread(self.image_path, cv2.IMREAD_COLOR) |
| h, w, _ = img.shape |
| person_center, s = self._box2cs([0, 0, w - 1, h - 1]) |
| r = 0 |
| trans = get_affine_transform(person_center, s, r, self.input_size) |
| input = cv2.warpAffine( |
| img, |
| trans, |
| (int(self.input_size[1]), int(self.input_size[0])), |
| flags=cv2.INTER_LINEAR, |
| borderMode=cv2.BORDER_CONSTANT, |
| borderValue=(0, 0, 0)) |
| input = self.transform(input) |
| meta = { |
| 'name': self.img_name, |
| 'center': person_center, |
| 'height': h, |
| 'width': w, |
| 'scale': s, |
| 'rotation': r |
| } |
| return input, meta |
|
|
|
|
| dataset_settings = { |
| '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'] |
| } |
| } |
|
|
| def get_palette(num_cls): |
| n = 18 |
| palette = [0] * (n * 3) |
| j = num_cls |
| lab = num_cls |
| palette[j * 3 + 0] = 0 |
| palette[j * 3 + 1] = 0 |
| palette[j * 3 + 2] = 0 |
| i = 0 |
| while lab: |
| palette[j * 3 + 0] = 255 |
| palette[j * 3 + 1] = 255 |
| palette[j * 3 + 2] = 255 |
| i += 1 |
| lab >>= 3 |
| return palette |
|
|
|
|
| def masking(image_path, class_num=0): |
| num_classes = dataset_settings['atr']['num_classes'] |
| input_size = dataset_settings['atr']['input_size'] |
| label = dataset_settings['atr']['label'] |
| print("Evaluating total class number {} with {}".format(num_classes, label)) |
|
|
| model = networks.init_model('resnet101', num_classes=num_classes, pretrained=None) |
| state_dict = torch.load('./ckpts/exp-schp-201908301523-atr.pth')['state_dict'] |
|
|
| from collections import OrderedDict |
| new_state_dict = OrderedDict() |
| for k, v in state_dict.items(): |
| name = k[7:] |
| new_state_dict[name] = v |
| model.load_state_dict(new_state_dict) |
| model.cuda() |
| model.eval() |
|
|
| transform = transforms.Compose([ |
| transforms.ToTensor(), |
| transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229]) |
| ]) |
| dataset = SimpleFileDataset(image_path=image_path, input_size=input_size, transform=transform) |
| dataloader = DataLoader(dataset) |
|
|
| if not os.path.exists('./outputs'): |
| os.makedirs('./outputs') |
|
|
| palette = get_palette(class_num) |
| 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.cuda()) |
| 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) |
|
|
| 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) |
| parsing_result_path = os.path.join('./outputs', img_name[:-4] + '.png') |
| output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8)) |
| output_img.putpalette(palette) |
| output_img.save(parsing_result_path) |
| gray_img = output_img.convert('L') |
|
|
| return gray_img |
|
|
