| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | import torchvision.transforms as transforms |
| | import cv2 |
| | import numpy as np |
| |
|
| | from .model import BiSeNet |
| |
|
| | mask_regions = { |
| | "Background":0, |
| | "Skin":1, |
| | "L-Eyebrow":2, |
| | "R-Eyebrow":3, |
| | "L-Eye":4, |
| | "R-Eye":5, |
| | "Eye-G":6, |
| | "L-Ear":7, |
| | "R-Ear":8, |
| | "Ear-R":9, |
| | "Nose":10, |
| | "Mouth":11, |
| | "U-Lip":12, |
| | "L-Lip":13, |
| | "Neck":14, |
| | "Neck-L":15, |
| | "Cloth":16, |
| | "Hair":17, |
| | "Hat":18 |
| | } |
| |
|
| | |
| | |
| | class SoftErosion(nn.Module): |
| | def __init__(self, kernel_size=15, threshold=0.6, iterations=1): |
| | super(SoftErosion, self).__init__() |
| | r = kernel_size // 2 |
| | self.padding = r |
| | self.iterations = iterations |
| | self.threshold = threshold |
| |
|
| | |
| | y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size)) |
| | dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2) |
| | kernel = dist.max() - dist |
| | kernel /= kernel.sum() |
| | kernel = kernel.view(1, 1, *kernel.shape) |
| | self.register_buffer('weight', kernel) |
| |
|
| | def forward(self, x): |
| | x = x.float() |
| | for i in range(self.iterations - 1): |
| | x = torch.min(x, F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding)) |
| | x = F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding) |
| |
|
| | mask = x >= self.threshold |
| | x[mask] = 1.0 |
| | x[~mask] /= x[~mask].max() |
| |
|
| | return x, mask |
| |
|
| | device = "cpu" |
| |
|
| | def init_parser(pth_path, mode="cpu"): |
| | global device |
| | device = mode |
| | n_classes = 19 |
| | net = BiSeNet(n_classes=n_classes) |
| | if device == "cuda": |
| | net.cuda() |
| | net.load_state_dict(torch.load(pth_path)) |
| | else: |
| | net.load_state_dict(torch.load(pth_path, map_location=torch.device('cpu'))) |
| | net.eval() |
| | return net |
| |
|
| |
|
| | def image_to_parsing(img, net): |
| | img = cv2.resize(img, (512, 512)) |
| | img = img[:,:,::-1] |
| | transform = transforms.Compose([ |
| | transforms.ToTensor(), |
| | transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) |
| | ]) |
| | img = transform(img.copy()) |
| | img = torch.unsqueeze(img, 0) |
| |
|
| | with torch.no_grad(): |
| | img = img.to(device) |
| | out = net(img)[0] |
| | parsing = out.squeeze(0).cpu().numpy().argmax(0) |
| | return parsing |
| |
|
| |
|
| | def get_mask(parsing, classes): |
| | res = parsing == classes[0] |
| | for val in classes[1:]: |
| | res += parsing == val |
| | return res |
| |
|
| |
|
| | def swap_regions(source, target, net, smooth_mask, includes=[1,2,3,4,5,10,11,12,13], blur=10): |
| | parsing = image_to_parsing(source, net) |
| |
|
| | if len(includes) == 0: |
| | return source, np.zeros_like(source) |
| |
|
| | include_mask = get_mask(parsing, includes) |
| | mask = np.repeat(include_mask[:, :, np.newaxis], 3, axis=2).astype("float32") |
| |
|
| | if smooth_mask is not None: |
| | mask_tensor = torch.from_numpy(mask.copy().transpose((2, 0, 1))).float().to(device) |
| | face_mask_tensor = mask_tensor[0] + mask_tensor[1] |
| | soft_face_mask_tensor, _ = smooth_mask(face_mask_tensor.unsqueeze_(0).unsqueeze_(0)) |
| | soft_face_mask_tensor.squeeze_() |
| | mask = np.repeat(soft_face_mask_tensor.cpu().numpy()[:, :, np.newaxis], 3, axis=2) |
| |
|
| | if blur > 0: |
| | mask = cv2.GaussianBlur(mask, (0, 0), blur) |
| |
|
| | resized_source = cv2.resize((source).astype("float32"), (512, 512)) |
| | resized_target = cv2.resize((target).astype("float32"), (512, 512)) |
| | result = mask * resized_source + (1 - mask) * resized_target |
| | result = cv2.resize(result.astype("uint8"), (source.shape[1], source.shape[0])) |
| |
|
| | return result |
| |
|
| | def mask_regions_to_list(values): |
| | out_ids = [] |
| | for value in values: |
| | if value in mask_regions.keys(): |
| | out_ids.append(mask_regions.get(value)) |
| | return out_ids |
| |
|
