262 lines
11 KiB
Python
262 lines
11 KiB
Python
# This file is not used by the tracking application and currently outdated
|
|
import torch
|
|
import torch.nn as nn
|
|
import geffnet.mobilenetv3 # geffnet.mobilenetv3._gen_mobilenet_v3 needs to be patched to return the parameters instead of instantiating the network
|
|
from geffnet.efficientnet_builder import round_channels
|
|
|
|
class DSConv2d(nn.Module):
|
|
def __init__(self, in_planes, out_planes, kernels_per_layer=4, groups=1, old=0):
|
|
super(DSConv2d, self).__init__()
|
|
if old == 2:
|
|
self.conv = nn.Sequential(
|
|
nn.Conv2d(in_planes, in_planes * kernels_per_layer, kernel_size=3, padding=1, groups=in_planes),
|
|
nn.Conv2d(in_planes * kernels_per_layer, out_planes, kernel_size=1, groups=groups)
|
|
)
|
|
elif old == 1:
|
|
self.conv = nn.Sequential(
|
|
nn.Conv2d(in_planes, in_planes * kernels_per_layer, kernel_size=3, padding=1, groups=in_planes, bias=False),
|
|
nn.BatchNorm2d(in_planes * kernels_per_layer),
|
|
nn.Conv2d(in_planes * kernels_per_layer, out_planes, kernel_size=1, groups=groups, bias=False),
|
|
nn.BatchNorm2d(out_planes),
|
|
nn.ReLU6(inplace=True)
|
|
)
|
|
else:
|
|
self.conv = nn.Sequential(
|
|
nn.Conv2d(in_planes, in_planes * kernels_per_layer, kernel_size=3, padding=1, groups=in_planes, bias=False),
|
|
nn.BatchNorm2d(in_planes * kernels_per_layer),
|
|
nn.ReLU6(inplace=True),
|
|
nn.Conv2d(in_planes * kernels_per_layer, out_planes, kernel_size=1, groups=groups, bias=False),
|
|
nn.BatchNorm2d(out_planes),
|
|
nn.ReLU6(inplace=True)
|
|
)
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
return x
|
|
|
|
class UNetUp(nn.Module):
|
|
def __init__(self, in_channels, residual_in_channels, out_channels, size, old=0):
|
|
super(UNetUp, self).__init__()
|
|
self.up = nn.Upsample(size=size, mode='bilinear', align_corners=True)
|
|
self.conv = DSConv2d(in_channels + residual_in_channels, out_channels, 1, 1, old=old)
|
|
def forward(self, x1, x2):
|
|
x1 = self.up(x1)
|
|
x = torch.cat([x2, x1], dim=1)
|
|
x = self.conv(x)
|
|
return x
|
|
|
|
# This is the gaze tracking model
|
|
class OpenSeeFaceGaze(geffnet.mobilenetv3.MobileNetV3):
|
|
def __init__(self):
|
|
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3(['small'])
|
|
super(OpenSeeFaceGaze, self).__init__(**kwargs)
|
|
self.up1 = UNetUp(576, 48, 64, (2,2), old=2)
|
|
self.up2 = UNetUp(64, 24, 32, (4,4), old=2)
|
|
self.up3 = UNetUp(32, 16, 15, (8,8), old=2)
|
|
self.group = DSConv2d(15, 3, kernels_per_layer=4, groups=3, old=2)
|
|
def _forward_impl(self, x):
|
|
x = self.conv_stem(x)
|
|
x = self.bn1(x)
|
|
x = self.act1(x)
|
|
r1 = None
|
|
r2 = None
|
|
r3 = None
|
|
for i, feature in enumerate(self.blocks):
|
|
x = feature(x)
|
|
if i == 3:
|
|
r3 = x
|
|
if i == 1:
|
|
r2 = x
|
|
if i == 0:
|
|
r1 = x
|
|
x = self.up1(x, r3)
|
|
x = self.up2(x, r2)
|
|
x = self.up3(x, r1)
|
|
x = self.group(x)
|
|
return x
|
|
def forward(self, x):
|
|
return self._forward_impl(x)
|
|
|
|
# This is the face detection model. Because the landmark model is very robust, it gets away with predicting very rough bounding boxes. It is fully convolutional and can be made to run on different resolutions. It was trained on 224x224 crops and the most reasonable results can be found in the range of 224x224 to 640x640.
|
|
class OpenSeeFaceDetect(geffnet.mobilenetv3.MobileNetV3):
|
|
def __init__(self, size="large", channel_multiplier=0.1):
|
|
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
|
|
super(OpenSeeFaceDetect, self).__init__(**kwargs)
|
|
if size == "large":
|
|
self.up1 = UNetUp(round_channels(960, channel_multiplier), round_channels(112, channel_multiplier), 256, (14,14), old=1)
|
|
self.up2 = UNetUp(256, round_channels(40, channel_multiplier), 128, (28,28), old=1)
|
|
self.up3 = UNetUp(128, round_channels(24, channel_multiplier), 64, (56,56), old=1)
|
|
self.group = DSConv2d(64, 2, kernels_per_layer=4, groups=2, old=1)
|
|
self.r1_i = 1
|
|
self.r2_i = 2
|
|
self.r3_i = 4
|
|
elif size == "small":
|
|
self.up1 = UNetUp(round_channels(576, channel_multiplier), round_channels(40, channel_multiplier), 256, (14,14), old=1)
|
|
self.up2 = UNetUp(256, round_channels(24, channel_multiplier), 128, (28,28), old=1)
|
|
self.up3 = UNetUp(128, round_channels(16, channel_multiplier), 64, (56,56), old=1)
|
|
self.group = DSConv2d(64, 2, kernels_per_layer=4, groups=2, old=1)
|
|
self.r1_i = 0
|
|
self.r2_i = 1
|
|
self.r3_i = 2
|
|
self.maxpool = nn.MaxPool2d(kernel_size=3, dilation=1, stride=1, padding=1)
|
|
def _forward_impl(self, x):
|
|
x = self.conv_stem(x)
|
|
x = self.bn1(x)
|
|
x = self.act1(x)
|
|
r2 = None
|
|
r3 = None
|
|
for i, feature in enumerate(self.blocks):
|
|
x = feature(x)
|
|
if i == self.r3_i:
|
|
r3 = x
|
|
if i == self.r2_i:
|
|
r2 = x
|
|
if i == self.r1_i:
|
|
r1 = x
|
|
x = self.up1(x, r3)
|
|
x = self.up2(x, r2)
|
|
x = self.up3(x, r1)
|
|
x = self.group(x)
|
|
x2 = self.maxpool(x)
|
|
return x, x2
|
|
def forward(self, x):
|
|
return self._forward_impl(x)
|
|
|
|
# Landmark detection model
|
|
# Models:
|
|
# 0: "small", 0.5
|
|
# 1: "small", 1.0
|
|
# 2: "large", 0.75
|
|
# 3: "large", 1.0
|
|
class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
|
|
def __init__(self, size="large", channel_multiplier=1.0):
|
|
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
|
|
super(OpenSeeFaceLandmarks, self).__init__(**kwargs)
|
|
if size == "large":
|
|
self.up1 = UNetUp(round_channels(960, channel_multiplier), round_channels(112, channel_multiplier), 256, (14,14))
|
|
self.up2 = UNetUp(256, round_channels(40, channel_multiplier), 198 * 1, (28,28))
|
|
self.group = DSConv2d(198 * 1, 198, kernels_per_layer=4, groups=3)
|
|
self.r2_i = 2
|
|
self.r3_i = 4
|
|
elif size == "small":
|
|
self.up1 = UNetUp(round_channels(576, channel_multiplier), round_channels(40, channel_multiplier), 256, (14,14))
|
|
self.up2 = UNetUp(256, round_channels(24, channel_multiplier), 198 * 1, (28,28))
|
|
self.group = DSConv2d(198 * 1, 198, kernels_per_layer=4, groups=3)
|
|
self.r2_i = 1
|
|
self.r3_i = 2
|
|
def _forward_impl(self, x):
|
|
x = self.conv_stem(x)
|
|
x = self.bn1(x)
|
|
x = self.act1(x)
|
|
r2 = None
|
|
r3 = None
|
|
for i, feature in enumerate(self.blocks):
|
|
x = feature(x)
|
|
if i == self.r3_i:
|
|
r3 = x
|
|
if i == self.r2_i:
|
|
r2 = x
|
|
x = self.up1(x, r3)
|
|
x = self.up2(x, r2)
|
|
x = self.group(x)
|
|
|
|
t_main = x[0:66].reshape((66, 28*28))
|
|
t_m = t_main.argmax(dim=1)
|
|
indices = t_m.unsqueeze(1)
|
|
t_conf = t_main.gather(1, indices).squeeze(1)
|
|
t_off_x = x[66:132].reshape((66, 28*28)).gather(1, indices).squeeze(1)
|
|
t_off_y = x[132:198].reshape((66, 28*28)).gather(1, indices).squeeze(1)
|
|
t_off_x = (223. * logit_arr(t_off_x) + 0.5).floor()
|
|
t_off_y = (223. * logit_arr(t_off_y) + 0.5).floor()
|
|
t_x = 223. * (t_m / 28.).floor() / 27. + t_off_x
|
|
t_y = 223. * t_m.remainder(28.).float() / 27. + t_off_y
|
|
x = (t_conf.mean(), torch.stack([t_x, t_y, t_conf], 1))
|
|
|
|
return x
|
|
def forward(self, x):
|
|
return self._forward_impl(x)
|
|
|
|
# Adaptive Wing Loss with offset layers and emphasis for eyes and eyebrows with 66 landmark points
|
|
def AdapWingLoss(pre_hm, gt_hm):
|
|
# pre_hm = pre_hm.to('cpu')
|
|
# gt_hm = gt_hm.to('cpu')
|
|
theta = 0.5
|
|
alpha = 2.1
|
|
w = 14
|
|
e = 1
|
|
A = w * (1 / (1 + torch.pow(theta / e, alpha - gt_hm))) * (alpha - gt_hm) * torch.pow(theta / e, alpha - gt_hm - 1) * (1 / e)
|
|
C = (theta * A - w * torch.log(1 + torch.pow(theta / e, alpha - gt_hm)))
|
|
|
|
batch_size = gt_hm.size()[0]
|
|
hm_num = gt_hm.size()[1]
|
|
|
|
mask = torch.zeros_like(gt_hm)
|
|
# W = 10
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
first = True
|
|
first_mask = None
|
|
for i in range(batch_size):
|
|
img_list = []
|
|
for j in range(hm_num // 3):
|
|
hm = np.round(gt_hm[i][j].cpu().numpy() * 255)
|
|
img_list.append(hm)
|
|
img_merge = cv2.merge(img_list)
|
|
img_dilate = cv2.morphologyEx(img_merge, cv2.MORPH_DILATE, kernel)
|
|
img_dilate[img_dilate < 51] = 1 # 0*W+1
|
|
img_dilate[img_dilate >= 51] = 11 # 1*W+1
|
|
img_dilate = np.array(img_dilate, dtype=np.int)
|
|
img_dilate = img_dilate.transpose(2, 0, 1)
|
|
dilated = torch.from_numpy(img_dilate).float()
|
|
if first:
|
|
first_mask = dilated
|
|
first = False
|
|
# These weights varied between training runs and model sizes
|
|
dilated[17:27][dilated[17:27] > 1] *= 1.4
|
|
dilated[17, dilated[17] > 1] *= 1.6
|
|
dilated[18, dilated[18] > 1] *= 1.8
|
|
dilated[25, dilated[25] > 1] *= 1.8
|
|
dilated[26, dilated[26] > 1] *= 1.6
|
|
dilated[36:48][dilated[36:48] > 1] *= 2.8
|
|
# Used for a very small model
|
|
#dilated[[37,38,40,41,43,44,46,47]][dilated[[37,38,40,41,43,44,46,47]] > 1] *= 20.8
|
|
mask[i] = torch.cat([dilated, dilated, dilated], 0)
|
|
|
|
diff_hm = torch.abs(gt_hm - pre_hm)
|
|
AWingLoss = A * diff_hm - C
|
|
idx = diff_hm < theta
|
|
AWingLoss[idx] = w * torch.log(1 + torch.pow(diff_hm / e, alpha - gt_hm))[idx]
|
|
|
|
AWingLoss *= mask
|
|
sum_loss = torch.sum(AWingLoss)
|
|
all_pixel = torch.sum(mask)
|
|
mean_loss = sum_loss / all_pixel
|
|
|
|
return first_mask.detach(), mean_loss
|
|
|
|
|
|
|
|
# Checkpoint test
|
|
if __name__== "__main__":
|
|
print("Checking gaze model")
|
|
m=OpenSeeFaceGaze()
|
|
ckpt = torch.load("gaze.pth")
|
|
m.load_state_dict(ckpt)
|
|
print("Checking detection model")
|
|
m=OpenSeeFaceDetect()
|
|
ckpt = torch.load("detection.pth")
|
|
m.load_state_dict(ckpt)
|
|
print("Checking lm_model0 model")
|
|
m=OpenSeeFaceLandmarks("small", 0.5)
|
|
ckpt = torch.load("lm_model0.pth")
|
|
m.load_state_dict(ckpt)
|
|
print("Checking lm_model1 model")
|
|
m=OpenSeeFaceLandmarks("small", 1.0)
|
|
ckpt = torch.load("lm_model1.pth")
|
|
m.load_state_dict(ckpt)
|
|
print("Checking lm_model2 model")
|
|
m=OpenSeeFaceLandmarks("large", 0.75)
|
|
ckpt = torch.load("lm_model2.pth")
|
|
m.load_state_dict(ckpt)
|
|
print("Checking lm_model3 model")
|
|
m=OpenSeeFaceLandmarks("large", 1.0)
|
|
ckpt = torch.load("lm_model3.pth")
|
|
m.load_state_dict(ckpt)
|