From 0f080e60fc9698eefb0bca7ce3cb05787a8f1767 Mon Sep 17 00:00:00 2001 From: David Chatel Date: Thu, 29 Feb 2024 14:28:19 +0100 Subject: [PATCH] initial --- .gitignore | 1 + BiSeNet/__init__.py | 283 ++++++++++++++++++++++++++++++++++++++++++++ BiSeNet/resnet.py | 109 +++++++++++++++++ __init__.py | 3 + nodes.py | 223 ++++++++++++++++++++++++++++++++++ utils.py | 168 ++++++++++++++++++++++++++ 6 files changed, 787 insertions(+) create mode 100644 .gitignore create mode 100644 BiSeNet/__init__.py create mode 100644 BiSeNet/resnet.py create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/BiSeNet/__init__.py b/BiSeNet/__init__.py new file mode 100644 index 0000000..5119e75 --- /dev/null +++ b/BiSeNet/__init__.py @@ -0,0 +1,283 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision + +from .resnet import Resnet18 +# from modules.bn import InPlaceABNSync as BatchNorm2d + + +class ConvBNReLU(nn.Module): + def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs): + super(ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_chan, + out_chan, + kernel_size = ks, + stride = stride, + padding = padding, + bias = False) + self.bn = nn.BatchNorm2d(out_chan) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = F.relu(self.bn(x)) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + +class BiSeNetOutput(nn.Module): + def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs): + super(BiSeNetOutput, self).__init__() + self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1) + self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = self.conv_out(x) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class AttentionRefinementModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(AttentionRefinementModule, self).__init__() + self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1) + self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False) + self.bn_atten = nn.BatchNorm2d(out_chan) + self.sigmoid_atten = nn.Sigmoid() + self.init_weight() + + def forward(self, x): + feat = self.conv(x) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv_atten(atten) + atten = self.bn_atten(atten) + atten = self.sigmoid_atten(atten) + out = torch.mul(feat, atten) + return out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + +class ContextPath(nn.Module): + def __init__(self, *args, **kwargs): + super(ContextPath, self).__init__() + self.resnet = Resnet18() + self.arm16 = AttentionRefinementModule(256, 128) + self.arm32 = AttentionRefinementModule(512, 128) + self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0) + + self.init_weight() + + def forward(self, x): + H0, W0 = x.size()[2:] + feat8, feat16, feat32 = self.resnet(x) + H8, W8 = feat8.size()[2:] + H16, W16 = feat16.size()[2:] + H32, W32 = feat32.size()[2:] + + avg = F.avg_pool2d(feat32, feat32.size()[2:]) + avg = self.conv_avg(avg) + avg_up = F.interpolate(avg, (H32, W32), mode='nearest') + + feat32_arm = self.arm32(feat32) + feat32_sum = feat32_arm + avg_up + feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest') + feat32_up = self.conv_head32(feat32_up) + + feat16_arm = self.arm16(feat16) + feat16_sum = feat16_arm + feat32_up + feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest') + feat16_up = self.conv_head16(feat16_up) + + return feat8, feat16_up, feat32_up # x8, x8, x16 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +### This is not used, since I replace this with the resnet feature with the same size +class SpatialPath(nn.Module): + def __init__(self, *args, **kwargs): + super(SpatialPath, self).__init__() + self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3) + self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0) + self.init_weight() + + def forward(self, x): + feat = self.conv1(x) + feat = self.conv2(feat) + feat = self.conv3(feat) + feat = self.conv_out(feat) + return feat + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class FeatureFusionModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(FeatureFusionModule, self).__init__() + self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0) + self.conv1 = nn.Conv2d(out_chan, + out_chan//4, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.conv2 = nn.Conv2d(out_chan//4, + out_chan, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + self.init_weight() + + def forward(self, fsp, fcp): + fcat = torch.cat([fsp, fcp], dim=1) + feat = self.convblk(fcat) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv1(atten) + atten = self.relu(atten) + atten = self.conv2(atten) + atten = self.sigmoid(atten) + feat_atten = torch.mul(feat, atten) + feat_out = feat_atten + feat + return feat_out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class BiSeNet(nn.Module): + def __init__(self, n_classes, *args, **kwargs): + super(BiSeNet, self).__init__() + self.cp = ContextPath() + ## here self.sp is deleted + self.ffm = FeatureFusionModule(256, 256) + self.conv_out = BiSeNetOutput(256, 256, n_classes) + self.conv_out16 = BiSeNetOutput(128, 64, n_classes) + self.conv_out32 = BiSeNetOutput(128, 64, n_classes) + self.init_weight() + + def forward(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + + feat_out = self.conv_out(feat_fuse) + feat_out16 = self.conv_out16(feat_cp8) + feat_out32 = self.conv_out32(feat_cp16) + + feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True) + feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True) + feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True) + return feat_out, feat_out16, feat_out32 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], [] + for name, child in self.named_children(): + child_wd_params, child_nowd_params = child.get_params() + if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput): + lr_mul_wd_params += child_wd_params + lr_mul_nowd_params += child_nowd_params + else: + wd_params += child_wd_params + nowd_params += child_nowd_params + return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params + + +if __name__ == "__main__": + net = BiSeNet(19) + net.cuda() + net.eval() + in_ten = torch.randn(16, 3, 640, 480).cuda() + out, out16, out32 = net(in_ten) + print(out.shape) + + net.get_params() diff --git a/BiSeNet/resnet.py b/BiSeNet/resnet.py new file mode 100644 index 0000000..aa2bf95 --- /dev/null +++ b/BiSeNet/resnet.py @@ -0,0 +1,109 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.model_zoo as modelzoo + +# from modules.bn import InPlaceABNSync as BatchNorm2d + +resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth' + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + def __init__(self, in_chan, out_chan, stride=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(in_chan, out_chan, stride) + self.bn1 = nn.BatchNorm2d(out_chan) + self.conv2 = conv3x3(out_chan, out_chan) + self.bn2 = nn.BatchNorm2d(out_chan) + self.relu = nn.ReLU(inplace=True) + self.downsample = None + if in_chan != out_chan or stride != 1: + self.downsample = nn.Sequential( + nn.Conv2d(in_chan, out_chan, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(out_chan), + ) + + def forward(self, x): + residual = self.conv1(x) + residual = F.relu(self.bn1(residual)) + residual = self.conv2(residual) + residual = self.bn2(residual) + + shortcut = x + if self.downsample is not None: + shortcut = self.downsample(x) + + out = shortcut + residual + out = self.relu(out) + return out + + +def create_layer_basic(in_chan, out_chan, bnum, stride=1): + layers = [BasicBlock(in_chan, out_chan, stride=stride)] + for i in range(bnum-1): + layers.append(BasicBlock(out_chan, out_chan, stride=1)) + return nn.Sequential(*layers) + + +class Resnet18(nn.Module): + def __init__(self): + super(Resnet18, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1) + self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2) + self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2) + self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2) + self.init_weight() + + def forward(self, x): + x = self.conv1(x) + x = F.relu(self.bn1(x)) + x = self.maxpool(x) + + x = self.layer1(x) + feat8 = self.layer2(x) # 1/8 + feat16 = self.layer3(feat8) # 1/16 + feat32 = self.layer4(feat16) # 1/32 + return feat8, feat16, feat32 + + def init_weight(self): + state_dict = modelzoo.load_url(resnet18_url) + self_state_dict = self.state_dict() + for k, v in state_dict.items(): + if 'fc' in k: continue + self_state_dict.update({k: v}) + self.load_state_dict(self_state_dict) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +if __name__ == "__main__": + net = Resnet18() + x = torch.randn(16, 3, 224, 224) + out = net(x) + print(out[0].size()) + print(out[1].size()) + print(out[2].size()) + net.get_params() diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3d998ea --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..168c5b7 --- /dev/null +++ b/nodes.py @@ -0,0 +1,223 @@ +import torch +from collections import defaultdict +from .utils import * + +class GenderFaceFilter: + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'faces': ('FACE',), + 'gender': (['male', 'female'],) + } + } + + RETURN_TYPES = ('FACE', 'FACE') + RETURN_NAMES = ('filtered', 'rest') + FUNCTION = 'run' + CATEGORY = 'facetools' + + def run(self, faces, gender): + gid = 0 if gender == 'female' else 1 + filtered = [] + rest = [] + for face in faces: + if face.gender == gid: + filtered.append(face) + else: + rest.append(face) + return (filtered, rest) + +class OrderedFaceFilter: + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'faces': ('FACE',), + 'criteria': (['area'],), + 'order': (['descending', 'ascending'],), + 'take_start': ('INT', {'default': 0, 'min': 0, 'step': 1}), + 'take_count': ('INT', {'default': 1, 'min': 1, 'step': 1}), + } + } + + RETURN_TYPES = ('FACE', 'FACE') + RETURN_NAMES = ('filtered', 'rest') + FUNCTION = 'run' + CATEGORY = 'facetools' + + def run(self, faces, criteria, order, take_start, take_count): + filtered = [] + rest = [] + funs = { + 'area': lambda face: face.w * face.h + } + sorted_faces = sorted(faces, key=funs[criteria], reverse=order == 'descending') + filtered = sorted_faces[take_start:take_start+take_count] + rest = sorted_faces[:take_start] + sorted_faces[take_start+take_count:] + return (filtered, rest) + +class FaceDetails: + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'faces': ('FACE',), + 'crop_size': ('INT', {'default': 512, 'min': 512, 'max': 1024, 'step': 128}), + 'crop_factor': ('FLOAT', {'default': 1.5, 'min': 1.0, 'max': 3, 'step': 0.1}), + 'mask_type': (mask_types,) + } + } + + RETURN_TYPES = ('IMAGE', 'MASK', 'WARP') + RETURN_NAMES = ('crops', 'masks', 'warps') + FUNCTION = 'run' + CATEGORY = 'facetools' + + def run(self, faces, crop_size, crop_factor, mask_type): + if len(faces) == 0: + empty_crop = torch.zeros((1,512,512,3)) + empty_mask = torch.zeros((1,512,512)) + empty_warp = np.array([ + [1,0,-512], + [0,1,-512], + ], dtype=np.float32) + return (empty_crop, empty_mask, [empty_warp]) + + crops = [] + masks = [] + warps = [] + for face in faces: + M, crop = crop_faces(face.image, face, crop_size, crop_factor) + mask = mask_crop(face, M, crop, mask_type) + crops.append(np.array(crop) / 255) + masks.append(np.array(mask)) + warps.append(M) + crops = torch.from_numpy(np.array(crops)) + masks = torch.from_numpy(np.array(masks)) + return (crops, masks, warps) + +class AlignFaces: + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'insightface': ('INSIGHTFACE',), + 'image': ('IMAGE',), + 'threshold': ('FLOAT', {'default': 0.5, 'min': 0.5, 'max': 1.0, 'step': 0.01}), + 'min_size': ('INT', {'default': 64, 'max': 512, 'step': 8}), + 'max_size': ('INT', {'default': 512, 'min': 512, 'step': 8}), + } + } + + RETURN_TYPES = ('FACE',) + RETURN_NAMES = ('faces',) + FUNCTION = 'run' + CATEGORY = 'facetools' + + def run(self, insightface, image, threshold, min_size, max_size): + faces = [] + images = (image * 255).type(torch.uint8).numpy() + for i, img in enumerate(images): + unfiltered_faces = get_faces(img, insightface) + for face in unfiltered_faces: + a, b, c, d = face.bbox + h = abs(d-b) + w = abs(c-a) + if face.det_score >= threshold and (h <= max_size or w <= max_size) and (min_size <= h or min_size <= w): + face.image_idx = i + face.image = img + faces.append(face) + return (faces,) + +class WarpFaceBack: + RETURN_TYPES = ('IMAGE',) + FUNCTION = 'run' + CATEGORY = 'facetools' + + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'images': ('IMAGE',), + 'face': ('FACE',), + 'crop': ('IMAGE',), + 'mask': ('MASK',), + 'warp': ('WARP',), + } + } + + def run(self, images, face, crop, mask, warp): + groups = defaultdict(list) + for f,c,m,w in zip(face, crop, mask, warp): + groups[f.image_idx].append((f.image,c,m,w)) + + results = [] + for i, image in enumerate(images): + if i not in groups: + result = image + else: + values = groups[i] + image = values[0][0] + crop, mask, warp = list(zip(*[x[1:] for x in values])) + warped_masks = [cv2.warpAffine(single_mask.numpy(), + cv2.invertAffineTransform(single_warp), + image.shape[1::-1]) + for single_warp, single_mask in zip(warp, mask)] + full_mask = np.add.reduce(warped_masks, axis=0)[...,None] + swapped = np.add.reduce([ + cv2.warpAffine(single_crop.numpy(), + cv2.invertAffineTransform(single_warp), + image.shape[1::-1] + ) * single_mask[..., None] + for single_crop, single_mask, single_warp in zip(crop, warped_masks, warp) + ], axis=0) / np.maximum(1, full_mask) + full_mask = np.minimum(1, full_mask) + result = (swapped + (1 - full_mask) * image / 255) + result = torch.from_numpy(result) + results.append(result) + + results = torch.stack(results) + return (results, ) + +class MergeWarps: + RETURN_TYPES = ('IMAGE','MASK','WARP') + FUNCTION = 'run' + CATEGORY = 'facetools' + + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'crop0': ('IMAGE',), + 'mask0': ('MASK',), + 'warp0': ('WARP',), + 'crop1': ('IMAGE',), + 'mask1': ('MASK',), + 'warp1': ('WARP',), + } + } + + def run(self, crop0, mask0, warp0, crop1, mask1, warp1): + crops = torch.vstack((crop0, crop1)) + masks = torch.vstack((mask0, mask1)) + warps = warp0 + warp1 + return (crops, masks, warps) + +NODE_CLASS_MAPPINGS = { + 'AlignFaces': AlignFaces, + 'WarpFacesBack': WarpFaceBack, + 'FaceDetails': FaceDetails, + 'MergeWarps': MergeWarps, + 'GenderFaceFilter': GenderFaceFilter, + 'OrderedFaceFilter': OrderedFaceFilter, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + 'AlignFaces': 'Align Faces', + 'WarpFacesBack': 'Warp Faces Back', + 'FaceDetails': 'Face Details', + 'MergeWarps': 'Merge Warps', + 'GenderFaceFilter': 'Gender Face Filter', + 'OrderedFaceFilter': 'Ordered Face Filter', +} diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..b659d1a --- /dev/null +++ b/utils.py @@ -0,0 +1,168 @@ +import os +import torch +import numpy as np +import cv2 +from insightface.utils import face_align +import mediapipe as mp +from scipy.spatial import ConvexHull +from folder_paths import models_dir +from .BiSeNet import BiSeNet + +def squareimg(img): + h, w, _ = img.shape + s = max(h, w) + ph, pw = (s - h) // 2, (s - w) // 2 + img = cv2.copyMakeBorder(img, ph, ph, pw, pw, cv2.BORDER_CONSTANT) + return (img, ph, pw) + +def transform_back(points, rot, pw, ph): + points = np.hstack((points, np.ones((len(points),1)))) @ rot.T + points -= np.array((pw, ph)) + return points + +def get_candidates(img, face_analysis): + simg, ph, pw = squareimg(img) + h, w, _ = img.shape + s = max(h, w) + candidates = [] + for i in range(4): + im = np.rot90(simg, i) + rot = cv2.getRotationMatrix2D((s/2, s/2), -90*i, 1) + dets = face_analysis.get(im) + for face in dets: + face.kps = transform_back(face.kps, rot, pw, ph) + face.landmark_2d_106 = transform_back(face.landmark_2d_106, rot, pw, ph) + + # original bbox and size + a,b,c,d = transform_back(face.bbox.reshape((2,2)), rot, pw, ph).flatten() + a,c = min(a,c),max(a,c) + b,d = min(b,d),max(b,d) + face.w, face.h = abs(c-a), abs(d-b) + + face.centroid = np.array(((a+c)/2, (b+d)/2)) + face.s = max(abs(a-c), abs(b-d)) + + face.angle = i*90 + candidates += dets + return candidates + +def group_faces(candidates): + groups = [] + for face in candidates: + group = None + for g in groups: + if any([np.linalg.norm(face.centroid - other.centroid) < np.sqrt(face.s)*2 for other in g]): + group = g + break + if group is None: + groups.append([face]) + else: + group.append(face) + return groups + +def get_best_faces(groups): + faces = [min(group, key=lambda face:abs(face.pose[2])) for group in groups] + return faces + +def get_faces(img, face_analysis): + candidates = get_candidates(img, face_analysis) + groups = group_faces(candidates) + faces = get_best_faces(groups) + return faces + +def crop_faces(img, face, crop_size, crop_factor): + M = face_align.estimate_norm(face.kps, crop_size, mode=None) + cx, cy = face.kps.mean(axis=0) + T1 = np.array([[1, 0, -cx], [0, 1, -cy], [0, 0, 1]]) + T2 = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]]) + S = np.array([[1/crop_factor, 0, 0], [0, 1/crop_factor, 0], [0, 0, 1]]) + M = M @ T2 @ S @ T1 + crop = cv2.warpAffine(img, M, (crop_size,crop_size)) + + return M, crop + +def get_face_mesh(crop): + with mp.solutions.face_mesh.FaceMesh(max_num_faces=10) as face_mesh: + mesh = face_mesh.process(crop) + h, w, _ = crop.shape + if mesh.multi_face_landmarks is not None: + all_pts = np.array([np.array([(w*l.x, h*l.y) for l in lmks.landmark]) for lmks in mesh.multi_face_landmarks], dtype=np.int32) + idx = np.argmin(np.abs(all_pts - np.array([w/2,h/2])).sum(axis=(1,2))) + points = all_pts[idx] + return points + else: + return None + +def mask_simple_square(face, M, crop): + # rotated bbox and size + a,b,c,d = face.bbox + rect = np.array([ + [a,b,1], + [a,d,1], + [c,b,1], + [c,d,1], + ]) @ M.T + lx, ly = [int(x) for x in np.min(rect, axis=0)] + hx, hy = [int(x) for x in np.max(rect, axis=0)] + mask = np.zeros((512,512), dtype=np.int32) + mask = cv2.rectangle(mask, (lx,ly), (hx,hy), 1, -1) + mask = mask.astype(np.float32) + return mask + +def mask_convex_hull(face, M, crop): + points = get_face_mesh(crop) + if points is None: return mask_simple_square(face, M, crop) + hull = ConvexHull(points) + mask = np.zeros((512,512), dtype=np.int32) + cv2.fillPoly(mask, [points[hull.vertices,:]], color=1) + mask = mask.astype(np.float32) + return mask + +def mask_BiSeNet(face, M, crop): + with torch.no_grad(): + bisenet = BiSeNet(n_classes=19) + bisenet.cuda() + model_path = os.path.join(models_dir, 'bisenet', '79999_iter.pth') + bisenet.load_state_dict(torch.load(model_path)) + bisenet.eval() + crop_t = torch.from_numpy(crop)[None].permute(0,3,1,2).contiguous().float().div(255).cuda() + segms_t = bisenet(crop_t)[0].argmax(1).float() + # convert segmentation to mask + # 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 + face_part_ids = torch.tensor([1,2,3,4,5,6,7,8,9,10,11,12,13]).cuda() + segms_t = torch.sum(segms_t.repeat(len(face_part_ids), 1,1,1) == face_part_ids[...,None,None,None], axis=0).float() + mask = segms_t[0].cpu() + return mask + +mask_types = [ + 'simple_square', + 'convex_hull', + 'BiSeNet', +] + +mask_funs = { + 'simple_square': mask_simple_square, + 'convex_hull': mask_convex_hull, + 'BiSeNet': mask_BiSeNet, +} + +def mask_crop(face, M, crop, mask_type): + mask = mask_funs[mask_type](face, M, crop) + return mask