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