| from model.SCHP import networks |
| from model.SCHP.utils.transforms import get_affine_transform, transform_logits |
|
|
| from collections import OrderedDict |
| import torch |
| import numpy as np |
| import cv2 |
| from PIL import Image |
| from torchvision import transforms |
|
|
| 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 |
|
|
| 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'], |
| } |
| } |
|
|
| class SCHP: |
| def __init__(self, ckpt_path, device): |
| dataset_type = None |
| if 'lip' in ckpt_path: |
| dataset_type = 'lip' |
| elif 'atr' in ckpt_path: |
| dataset_type = 'atr' |
| elif 'pascal' in ckpt_path: |
| dataset_type = 'pascal' |
| assert dataset_type is not None, 'Dataset type not found in checkpoint path' |
| self.device = device |
| self.num_classes = dataset_settings[dataset_type]['num_classes'] |
| self.input_size = dataset_settings[dataset_type]['input_size'] |
| self.aspect_ratio = self.input_size[1] * 1.0 / self.input_size[0] |
| self.palette = get_palette(self.num_classes) |
|
|
| self.label = dataset_settings[dataset_type]['label'] |
| self.model = networks.init_model('resnet101', num_classes=self.num_classes, pretrained=None).to(device) |
| self.load_ckpt(ckpt_path) |
| self.model.eval() |
| |
| self.transform = transforms.Compose([ |
| transforms.ToTensor(), |
| transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229]) |
| ]) |
| self.upsample = torch.nn.Upsample(size=self.input_size, mode='bilinear', align_corners=True) |
|
|
|
|
| def load_ckpt(self, ckpt_path): |
| rename_map = { |
| "decoder.conv3.2.weight": "decoder.conv3.3.weight", |
| "decoder.conv3.3.weight": "decoder.conv3.4.weight", |
| "decoder.conv3.3.bias": "decoder.conv3.4.bias", |
| "decoder.conv3.3.running_mean": "decoder.conv3.4.running_mean", |
| "decoder.conv3.3.running_var": "decoder.conv3.4.running_var", |
| "fushion.3.weight": "fushion.4.weight", |
| "fushion.3.bias": "fushion.4.bias", |
| } |
| state_dict = torch.load(ckpt_path, map_location='cpu')['state_dict'] |
| new_state_dict = OrderedDict() |
| for k, v in state_dict.items(): |
| name = k[7:] |
| new_state_dict[name] = v |
| new_state_dict_ = OrderedDict() |
| for k, v in list(new_state_dict.items()): |
| if k in rename_map: |
| new_state_dict_[rename_map[k]] = v |
| else: |
| new_state_dict_[k] = v |
| self.model.load_state_dict(new_state_dict_, strict=False) |
|
|
| 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 preprocess(self, image): |
| if isinstance(image, str): |
| img = cv2.imread(image, cv2.IMREAD_COLOR) |
| elif isinstance(image, Image.Image): |
| |
| img = np.array(image) |
| |
| 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).to(self.device).unsqueeze(0) |
| meta = { |
| 'center': person_center, |
| 'height': h, |
| 'width': w, |
| 'scale': s, |
| 'rotation': r |
| } |
| return input, meta |
|
|
|
|
| def __call__(self, image_or_path): |
| if isinstance(image_or_path, list): |
| image_list = [] |
| meta_list = [] |
| for image in image_or_path: |
| image, meta = self.preprocess(image) |
| image_list.append(image) |
| meta_list.append(meta) |
| image = torch.cat(image_list, dim=0) |
| else: |
| image, meta = self.preprocess(image_or_path) |
| meta_list = [meta] |
| |
| output = self.model(image) |
| |
| upsample_outputs = self.upsample(output) |
| upsample_outputs = upsample_outputs.permute(0, 2, 3, 1) |
|
|
| output_img_list = [] |
| for upsample_output, meta in zip(upsample_outputs, meta_list): |
| c, s, w, h = meta['center'], meta['scale'], meta['width'], meta['height'] |
| logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=self.input_size) |
| parsing_result = np.argmax(logits_result, axis=2) |
| output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8)) |
| output_img.putpalette(self.palette) |
| output_img_list.append(output_img) |
|
|
| return output_img_list[0] if len(output_img_list) == 1 else output_img_list |
