|
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 |
