import torch
from PIL import Image
from facexlib.parsing import init_parsing_model
from matplotlib import pyplot as plt
from torch import nn
from torchvision.transforms.functional import pil_to_tensor, normalize
class BiSeNetFaceXLibSegmenter(nn.Module):
"""
https://github.com/xinntao/facexlib/blob/master/inference/inference_parsing.py
https://github.com/xinntao/facexlib/blob/master/inference/inference_parsing_parsenet.py
"""
def __init__(self, model_name="bisenet", device="cuda"):
super().__init__()
assert model_name in ["bisenet", "parsenet"]
self.net = init_parsing_model(model_name=model_name, half=True, device=device)
# self.net.load_state_dict(torch.load("face_parsing/79999_iter.pth"))
self.mean = [0.485, 0.456, 0.406] if model_name == "bisenet" else [0.5, 0.5, 0.5]
self.std = [0.229, 0.224, 0.225] if model_name == "bisenet" else [0.5, 0.5, 0.5]
# attributions = [0 'bg', 1 'skin', 2 'l_brow', 3 'r_brow', 4 'l_eye', 5 'r_eye',
# 6 'eye_g', 7 'l_ear', 8 'r_ear', 9 'ear_r', 10 'nose',
# 11 'mouth', 12 'u_lip', 13 'l_lip', 14 'neck', 15 'neck_l',
# 16 'cloth', 17 'hair', 18 'hat']
def forward(self, images: torch.Tensor):
images = normalize(images, self.mean, self.std)
out = self.net(images)
category_masks = out[0].argmax(1)
face_skin_masks = torch.where(category_masks == 0,
torch.ones_like(category_masks),
torch.zeros_like(category_masks))
return face_skin_masks
