copilot/OpenSeeFace
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Changed landmark decoding in model.py to only be active for inference.

Browse source
This commit is contained in:
Emiliana 2020-08-10 12:08:54 +02:00 • committed by GitHub
parent cc7667734b
commit cc0ce01f09
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

4
model.py
View file

@ -128,7 +128,7 @@ class OpenSeeFaceDetect(geffnet.mobilenetv3.MobileNetV3):
# 2: "large", 0.75
# 3: "large", 1.0
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)
super(OpenSeeFaceLandmarks, self).__init__(**kwargs)
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.r2_i = 1
self.r3_i = 2
self.inference = inference
def _forward_impl(self, x):
x = self.conv_stem(x)
x = self.bn1(x)
@ -159,6 +160,7 @@ class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
x = self.up2(x, r2)
x = self.group(x)
if self.inference:
t_main = x[0:66].reshape((66, 28*28))
t_m = t_main.argmax(dim=1)
indices = t_m.unsqueeze(1)