|
|
'''
|
|
|
@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
|
|
|
@author: yangxy (yangtao9009@gmail.com)
|
|
|
'''
|
|
|
import os
|
|
|
import cv2
|
|
|
import torch
|
|
|
import numpy as np
|
|
|
from face_parse.parse_model import ParseNet
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
from face_parse.model import BiSeNet
|
|
|
import torchvision.transforms as transforms
|
|
|
|
|
|
class FaceParse(object):
|
|
|
def __init__(self, base_dir='./', model='ParseNet-latest', device='cuda', mask_map = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0]):
|
|
|
self.mfile = os.path.join(base_dir, model+'.pth')
|
|
|
self.size = 512
|
|
|
self.device = device
|
|
|
|
|
|
'''
|
|
|
0: 'background' 1: 'skin' 2: 'nose'
|
|
|
3: 'eye_g' 4: 'l_eye' 5: 'r_eye'
|
|
|
6: 'l_brow' 7: 'r_brow' 8: 'l_ear'
|
|
|
9: 'r_ear' 10: 'mouth' 11: 'u_lip'
|
|
|
12: 'l_lip' 13: 'hair' 14: 'hat'
|
|
|
15: 'ear_r' 16: 'neck_l' 17: 'neck'
|
|
|
18: 'cloth'
|
|
|
'''
|
|
|
|
|
|
|
|
|
|
|
|
self.MASK_COLORMAP = mask_map
|
|
|
|
|
|
self.load_model()
|
|
|
|
|
|
def load_model(self):
|
|
|
self.faceparse = ParseNet(self.size, self.size, 32, 64, 19, norm_type='bn', relu_type='LeakyReLU', ch_range=[32, 256])
|
|
|
self.faceparse.load_state_dict(torch.load(self.mfile, map_location=torch.device('cpu')))
|
|
|
self.faceparse.to(self.device)
|
|
|
self.faceparse.eval()
|
|
|
|
|
|
def process(self, im, masks=[0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0]):
|
|
|
im = cv2.resize(im, (self.size, self.size))
|
|
|
imt = self.img2tensor(im)
|
|
|
with torch.no_grad():
|
|
|
pred_mask, sr_img_tensor = self.faceparse(imt)
|
|
|
mask = self.tenor2mask(pred_mask, masks)
|
|
|
|
|
|
return mask
|
|
|
|
|
|
def process_tensor(self, imt):
|
|
|
imt = F.interpolate(imt.flip(1)*2-1, (self.size, self.size))
|
|
|
pred_mask, sr_img_tensor = self.faceparse(imt)
|
|
|
|
|
|
mask = pred_mask.argmax(dim=1)
|
|
|
for idx, color in enumerate(self.MASK_COLORMAP):
|
|
|
mask = torch.where(mask==idx, color, mask)
|
|
|
|
|
|
mask = mask.unsqueeze(0)
|
|
|
|
|
|
return mask
|
|
|
|
|
|
def img2tensor(self, img):
|
|
|
img = img[..., ::-1]
|
|
|
img = img / 255. * 2 - 1
|
|
|
img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(self.device)
|
|
|
return img_tensor.float()
|
|
|
|
|
|
def tenor2mask(self, tensor, masks):
|
|
|
if len(tensor.shape) < 4:
|
|
|
tensor = tensor.unsqueeze(0)
|
|
|
if tensor.shape[1] > 1:
|
|
|
tensor = tensor.argmax(dim=1)
|
|
|
|
|
|
tensor = tensor.squeeze(1).data.cpu().numpy()
|
|
|
color_maps = []
|
|
|
for t in tensor:
|
|
|
|
|
|
tmp_img = np.zeros(tensor.shape[1:])
|
|
|
for idx, color in enumerate(masks):
|
|
|
tmp_img[t == idx] = color
|
|
|
color_maps.append(tmp_img.astype(np.uint8))
|
|
|
return color_maps
|
|
|
|
|
|
|
|
|
|
|
|
class FaceParse_v2(object):
|
|
|
def __init__(self, device='cuda', mask_map = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0]):
|
|
|
self.mfile = '/apdcephfs/private_quincheng/Expression/face-parsing.PyTorch/res/cp/79999_iter.pth'
|
|
|
self.size = 512
|
|
|
self.device = device
|
|
|
|
|
|
'''
|
|
|
0: 'background' 1: 'skin' 2: 'nose'
|
|
|
3: 'eye_g' 4: 'l_eye' 5: 'r_eye'
|
|
|
6: 'l_brow' 7: 'r_brow' 8: 'l_ear'
|
|
|
9: 'r_ear' 10: 'mouth' 11: 'u_lip'
|
|
|
12: 'l_lip' 13: 'hair' 14: 'hat'
|
|
|
15: 'ear_r' 16: 'neck_l' 17: 'neck'
|
|
|
18: 'cloth'
|
|
|
'''
|
|
|
|
|
|
|
|
|
|
|
|
self.MASK_COLORMAP = mask_map
|
|
|
self.load_model()
|
|
|
self.to_tensor = transforms.Compose([
|
|
|
transforms.ToTensor(),
|
|
|
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
|
|
|
])
|
|
|
|
|
|
def load_model(self):
|
|
|
self.faceparse = BiSeNet(n_classes=19)
|
|
|
self.faceparse.load_state_dict(torch.load(self.mfile))
|
|
|
self.faceparse.to(self.device)
|
|
|
self.faceparse.eval()
|
|
|
|
|
|
def process(self, im, masks=[0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0]):
|
|
|
im = cv2.resize(im[...,::-1], (self.size, self.size))
|
|
|
im = self.to_tensor(im)
|
|
|
imt = torch.unsqueeze(im, 0).to(self.device)
|
|
|
with torch.no_grad():
|
|
|
pred_mask = self.faceparse(imt)[0]
|
|
|
mask = self.tenor2mask(pred_mask, masks)
|
|
|
return mask
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def tenor2mask(self, tensor, masks):
|
|
|
if len(tensor.shape) < 4:
|
|
|
tensor = tensor.unsqueeze(0)
|
|
|
if tensor.shape[1] > 1:
|
|
|
tensor = tensor.argmax(dim=1)
|
|
|
|
|
|
tensor = tensor.squeeze(1).data.cpu().numpy()
|
|
|
color_maps = []
|
|
|
for t in tensor:
|
|
|
|
|
|
tmp_img = np.zeros(tensor.shape[1:])
|
|
|
for idx, color in enumerate(masks):
|
|
|
tmp_img[t == idx] = color
|
|
|
color_maps.append(tmp_img.astype(np.uint8))
|
|
|
return color_maps |
