Spaces:
Sleeping
Sleeping
| from collections import OrderedDict | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from SCHP import networks | |
| from SCHP.utils.transforms import get_affine_transform, transform_logits | |
| 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:] # remove `module.` | |
| 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): | |
| # to cv2 format | |
| img = np.array(image) | |
| h, w, _ = img.shape | |
| # Get person center and scale | |
| 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[0][-1]) | |
| upsample_outputs = self.upsample(output) | |
| upsample_outputs = upsample_outputs.permute(0, 2, 3, 1) # BCHW -> BHWC | |
| 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 | |