initial
This commit is contained in:
David Chatel
@@ -0,0 +1 @@
|
||||
__pycache__
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -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',
|
||||
}
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user