Source code for pyVHR.extraction.skin_extraction_methods
from numba import cuda, njit, prange
import math
import numpy as np
import torchvision.transforms as transforms
import torch
from numba import prange, njit, cuda
from pyVHR.resources.faceparsing.model import BiSeNet
import os
import pyVHR
from scipy.spatial import ConvexHull
from PIL import Image, ImageDraw
import requests
"""
This module defines classes or methods used for skin extraction.
"""
### functions and parameters ###
[docs]class SkinProcessingParams():
"""
This class contains usefull parameters used by this module.
RGB_LOW_TH (numpy.int32): RGB low-threshold value.
RGB_HIGH_TH (numpy.int32): RGB high-threshold value.
"""
RGB_LOW_TH = np.int32(55)
RGB_HIGH_TH = np.int32(200)
[docs]def bbox2_CPU(img):
"""
Args:
img (ndarray): ndarray with shape [rows, columns, rgb_channels].
Returns:
Four cropping coordinates (row, row, column, column) for removing black borders (RGB [O,O,O]) from img.
"""
rows = np.any(img, axis=1)
cols = np.any(img, axis=0)
nzrows = np.nonzero(rows)
nzcols = np.nonzero(cols)
if nzrows[0].size == 0 or nzcols[0].size == 0:
return -1, -1, -1, -1
rmin, rmax = np.nonzero(rows)[0][[0, -1]]
cmin, cmax = np.nonzero(cols)[0][[0, -1]]
return rmin, rmax, cmin, cmax
### SKIN EXTRACTION CLASSES ###
[docs]class SkinExtractionFaceParsing():
"""
This class performs skin extraction on CPU/GPU using Face Parsing.
https://github.com/zllrunning/face-parsing.PyTorch
"""
def __init__(self):
"""
Args:
device (str): This class can execute code on 'CPU'.
"""
n_classes = 19
self.net = BiSeNet(n_classes=n_classes)
save_pth = os.path.dirname(
pyVHR.resources.faceparsing.model.__file__) + "/79999_iter.pth"
if not os.path.isfile(save_pth):
url = "https://github.com/phuselab/pyVHR/blob/main/pyVHR/resources/faceparsing/79999_iter.pth"
print('Downloading faceparsing model...')
r = requests.get(url, allow_redirects=True)
open(save_pth, 'wb').write(r.content)
self.net.load_state_dict(torch.load(save_pth, map_location=torch.device('cpu')))
self.net.eval()
self.to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
[docs] def extract_skin(self, image, ldmks):
"""
This method extract the skin from an image using Face Parsing.
Landmarks (ldmks) are used to create a facial bounding box for cropping the face; this way
the network used in Face Parsing is more accurate.
Args:
image (uint8 ndarray): ndarray with shape [rows, columns, rgb_channels].
ldmks (float32 ndarray): ndarray with shape [num_landmarks, xy_coordinates].
Returns:
Cropped skin-image and non-cropped skin-image; both are uint8 ndarray with shape [rows, columns, rgb_channels].
"""
# crop with bounding box of ldmks; the network works better if the bounding box is bigger
aviable_ldmks = ldmks[ldmks[:,0] >= 0][:,:2]
min_y, min_x = np.min(aviable_ldmks, axis=0)
max_y, max_x = np.max(aviable_ldmks, axis=0)
min_y *= 0.90
min_x *= 0.90
max_y = max_y * 1.10 if max_y * 1.10 < image.shape[0] else image.shape[0]
max_x = max_x * 1.10 if max_x * 1.10 < image.shape[1] else image.shape[1]
cropped_image = np.copy(image[int(min_y):int(max_y),int(min_x):int(max_x) ,:])
nda_im = np.array(cropped_image)
# prepare the image for the bisenet network
cropped_image = self.to_tensor(cropped_image)
cropped_image = torch.unsqueeze(cropped_image, 0)
cropped_skin_img = self.extraction(cropped_image, nda_im)
# recreate full image using cropped_skin_img
full_skin_image = np.zeros_like(image)
full_skin_image[int(min_y):int(max_y),int(min_x):int(max_x) ,:] = cropped_skin_img
return cropped_skin_img, full_skin_image
[docs] def extraction(self, im, nda_im):
"""
This method performs skin extraction using Face Parsing.
Args:
im (torch.Tensor): torch.Tensor with size [rows, columns, rgb_channels]
nda_im (uint8 ndarray): ndarray with shape [rows, columns, rgb_channels].
Returns:
skin-image as uint8 ndarray with shape [rows, columns, rgb_channels].
"""
with torch.no_grad():
### bisenet skin detection ###
out = self.net(im)[0]
### gpu cuda skin copy ###
parsing = out.squeeze(0).argmax(0).numpy()
parsing = parsing.astype(np.int32)
nda_im = nda_im.astype(np.uint8)
# kernel skin copy
return kernel_skin_copy_and_filter(nda_im, parsing, np.int32(SkinProcessingParams.RGB_LOW_TH), np.int32(SkinProcessingParams.RGB_HIGH_TH))
# SkinExtractionFaceParsing class kernels #
[docs]@njit('uint8[:,:,:](uint8[:,:,:], int32[:,:], int32, int32)', parallel=True, nogil=True)
def kernel_skin_copy_and_filter(orig, pars, RGB_LOW_TH, RGB_HIGH_TH):
"""
This method removes pixels from the image 'orig' that are not skin, or
that are outside the RGB range [RGB_LOW_TH, RGB_HIGH_TH] (extremes are included).
"""
new = np.zeros_like(orig)
for x in prange(orig.shape[0]):
for y in prange(orig.shape[1]):
# skin class = 1, nose = 10
if pars[x, y] == 1 or pars[x, y] == 10:
if not ((orig[x, y, 0] <= RGB_LOW_TH and orig[x, y, 1] <= RGB_LOW_TH and orig[x, y, 2] <= RGB_LOW_TH) or
(orig[x, y, 0] >= RGB_HIGH_TH and orig[x, y, 1] >= RGB_HIGH_TH and orig[x, y, 2] >= RGB_HIGH_TH)):
new[x, y, 0] = orig[x, y, 0]
new[x, y, 1] = orig[x, y, 1]
new[x, y, 2] = orig[x, y, 2]
return new
[docs]class SkinExtractionConvexHull:
"""
This class performs skin extraction on CPU/GPU using a Convex Hull segmentation obtained from facial landmarks.
"""
def __init__(self):
"""
Args:
device (str): This class can execute code on 'CPU'.
"""
[docs] def extract_skin(self,image, ldmks):
"""
This method extract the skin from an image using Convex Hull segmentation.
Args:
image (uint8 ndarray): ndarray with shape [rows, columns, rgb_channels].
ldmks (float32 ndarray): landmarks used to create the Convex Hull; ldmks is a ndarray with shape [num_landmarks, xy_coordinates].
Returns:
Cropped skin-image and non-cropped skin-image; both are uint8 ndarray with shape [rows, columns, rgb_channels].
"""
from pyVHR.extraction.sig_processing import MagicLandmarks
aviable_ldmks = ldmks[ldmks[:,0] >= 0][:,:2]
# face_mask convex hull
hull = ConvexHull(aviable_ldmks)
verts = [(aviable_ldmks[v,0], aviable_ldmks[v,1]) for v in hull.vertices]
img = Image.new('L', image.shape[:2], 0)
ImageDraw.Draw(img).polygon(verts, outline=1, fill=1)
mask = np.array(img)
mask = np.expand_dims(mask,axis=0).T
# left eye convex hull
left_eye_ldmks = ldmks[MagicLandmarks.left_eye]
aviable_ldmks = left_eye_ldmks[left_eye_ldmks[:,0] >= 0][:,:2]
if len(aviable_ldmks) > 3:
hull = ConvexHull(aviable_ldmks)
verts = [(aviable_ldmks[v,0], aviable_ldmks[v,1]) for v in hull.vertices]
img = Image.new('L', image.shape[:2], 0)
ImageDraw.Draw(img).polygon(verts, outline=1, fill=1)
left_eye_mask = np.array(img)
left_eye_mask = np.expand_dims(left_eye_mask,axis=0).T
else:
left_eye_mask = np.ones((image.shape[0], image.shape[1],1),dtype=np.uint8)
# right eye convex hull
right_eye_ldmks = ldmks[MagicLandmarks.right_eye]
aviable_ldmks = right_eye_ldmks[right_eye_ldmks[:,0] >= 0][:,:2]
if len(aviable_ldmks) > 3:
hull = ConvexHull(aviable_ldmks)
verts = [(aviable_ldmks[v,0], aviable_ldmks[v,1]) for v in hull.vertices]
img = Image.new('L', image.shape[:2], 0)
ImageDraw.Draw(img).polygon(verts, outline=1, fill=1)
right_eye_mask = np.array(img)
right_eye_mask = np.expand_dims(right_eye_mask,axis=0).T
else:
right_eye_mask = np.ones((image.shape[0], image.shape[1],1),dtype=np.uint8)
# mounth convex hull
mounth_ldmks = ldmks[MagicLandmarks.mounth]
aviable_ldmks = mounth_ldmks[mounth_ldmks[:,0] >= 0][:,:2]
if len(aviable_ldmks) > 3:
hull = ConvexHull(aviable_ldmks)
verts = [(aviable_ldmks[v,0], aviable_ldmks[v,1]) for v in hull.vertices]
img = Image.new('L', image.shape[:2], 0)
ImageDraw.Draw(img).polygon(verts, outline=1, fill=1)
mounth_mask = np.array(img)
mounth_mask = np.expand_dims(mounth_mask,axis=0).T
else:
mounth_mask = np.ones((image.shape[0], image.shape[1],1),dtype=np.uint8)
# apply masks and crop
skin_image = image * mask * (1-left_eye_mask) * (1-right_eye_mask) * (1-mounth_mask)
rmin, rmax, cmin, cmax = bbox2_CPU(skin_image)
cropped_skin_im = skin_image
if rmin >= 0 and rmax >= 0 and cmin >= 0 and cmax >= 0 and rmax-rmin >= 0 and cmax-cmin >= 0:
cropped_skin_im = skin_image[int(rmin):int(rmax), int(cmin):int(cmax)]
return cropped_skin_im, skin_image