From 2cca2cba5793173ef91a4ff7cfb705c460c45819 Mon Sep 17 00:00:00 2001 From: Emiliana <38952746+emilianavt@users.noreply.github.com> Date: Thu, 24 Sep 2020 17:57:35 +0200 Subject: [PATCH] Fixed the number of layers in 30 point model and added missing batch dimension (#7) --- model.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/model.py b/model.py index b3d0532..b9faffe 100644 --- a/model.py +++ b/model.py @@ -161,12 +161,12 @@ class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3): x = self.group(x) if self.inference: - t_main = x[0:66].reshape((66, 28*28)) + t_main = x[:, 0:66].reshape((-1, 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 = x[:, 66:132].reshape((-1, 66, 28*28)).gather(1, indices).squeeze(1) + t_off_y = x[:, 132:198].reshape((-1, 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 @@ -203,12 +203,12 @@ class OpenSeeFaceLandmarks30Pt(geffnet.mobilenetv3.MobileNetV3): x = self.group(x) if self.inference: - t_main = x[0:30].reshape((30, 7*7)) + t_main = x[:, 0:30].reshape((-1, 30, 7*7)) 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[30:60].reshape((60, 7*7)).gather(1, indices).squeeze(1) - t_off_y = x[60:90].reshape((60, 7*7)).gather(1, indices).squeeze(1) + t_off_x = x[:, 30:60].reshape((-1, 30, 7*7)).gather(1, indices).squeeze(1) + t_off_y = x[:, 60:90].reshape((-1, 30, 7*7)).gather(1, indices).squeeze(1) t_off_x = 55. * logit_arr(t_off_x) t_off_y = 55. * logit_arr(t_off_y) t_x = 55. * (t_m / 7.).floor() / 6. + t_off_x