Merge branch 'master' of github.com:emilianavt/OpenSeeFace

This commit is contained in:
Emiliana 2020-08-10 14:40:55 +02:00
commit 76c04f3be5

View file

@ -128,7 +128,7 @@ class OpenSeeFaceDetect(geffnet.mobilenetv3.MobileNetV3):
# 2: "large", 0.75 # 2: "large", 0.75
# 3: "large", 1.0 # 3: "large", 1.0
class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3): class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
def __init__(self, size="large", channel_multiplier=1.0): def __init__(self, size="large", channel_multiplier=1.0, inference=False):
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier) kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
super(OpenSeeFaceLandmarks, self).__init__(**kwargs) super(OpenSeeFaceLandmarks, self).__init__(**kwargs)
if size == "large": if size == "large":
@ -143,6 +143,7 @@ class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
self.group = DSConv2d(198 * 1, 198, kernels_per_layer=4, groups=3) self.group = DSConv2d(198 * 1, 198, kernels_per_layer=4, groups=3)
self.r2_i = 1 self.r2_i = 1
self.r3_i = 2 self.r3_i = 2
self.inference = inference
def _forward_impl(self, x): def _forward_impl(self, x):
x = self.conv_stem(x) x = self.conv_stem(x)
x = self.bn1(x) x = self.bn1(x)
@ -159,17 +160,18 @@ class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
x = self.up2(x, r2) x = self.up2(x, r2)
x = self.group(x) x = self.group(x)
t_main = x[0:66].reshape((66, 28*28)) if self.inference:
t_m = t_main.argmax(dim=1) t_main = x[0:66].reshape((66, 28*28))
indices = t_m.unsqueeze(1) t_m = t_main.argmax(dim=1)
t_conf = t_main.gather(1, indices).squeeze(1) indices = t_m.unsqueeze(1)
t_off_x = x[66:132].reshape((66, 28*28)).gather(1, indices).squeeze(1) t_conf = t_main.gather(1, indices).squeeze(1)
t_off_y = x[132:198].reshape((66, 28*28)).gather(1, indices).squeeze(1) t_off_x = x[66:132].reshape((66, 28*28)).gather(1, indices).squeeze(1)
t_off_x = (223. * logit_arr(t_off_x) + 0.5).floor() t_off_y = x[132:198].reshape((66, 28*28)).gather(1, indices).squeeze(1)
t_off_y = (223. * logit_arr(t_off_y) + 0.5).floor() t_off_x = (223. * logit_arr(t_off_x) + 0.5).floor()
t_x = 223. * (t_m / 28.).floor() / 27. + t_off_x t_off_y = (223. * logit_arr(t_off_y) + 0.5).floor()
t_y = 223. * t_m.remainder(28.).float() / 27. + t_off_y t_x = 223. * (t_m / 28.).floor() / 27. + t_off_x
x = (t_conf.mean(), torch.stack([t_x, t_y, t_conf], 1)) 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 return x
def forward(self, x): def forward(self, x):