From e405afdca83189e0b44b9305cbf819abad0f8454 Mon Sep 17 00:00:00 2001 From: Sacha Refshauge Date: Wed, 2 Oct 2019 15:44:02 +1000 Subject: [PATCH] Select only the important class dim for inference --- detect.py | 2 +- models/retinaface.py | 4 ++-- test.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/detect.py b/detect.py index a1492bee5..51bed6ae8 100755 --- a/detect.py +++ b/detect.py @@ -102,7 +102,7 @@ def load_model(model, pretrained_path, load_to_cpu): boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance']) boxes = boxes * scale / resize boxes = boxes.cpu().numpy() - scores = conf.squeeze(0).data.cpu().numpy()[:, 1] + scores = conf.squeeze(0).data.cpu().numpy() landms = decode_landm(landms.data.squeeze(0), prior_data, cfg['variance']) scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2], diff --git a/models/retinaface.py b/models/retinaface.py index 246c1ab46..db91a8cba 100644 --- a/models/retinaface.py +++ b/models/retinaface.py @@ -121,5 +121,5 @@ def forward(self,inputs): if self.phase == 'train': output = (bbox_regressions, classifications, ldm_regressions) else: - output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions) - return output \ No newline at end of file + output = (bbox_regressions, F.softmax(classifications, dim=-1).select(2, 1), ldm_regressions) + return output diff --git a/test.py b/test.py index e4a988595..f8d46d3ab 100755 --- a/test.py +++ b/test.py @@ -121,7 +121,7 @@ def load_model(model, pretrained_path, load_to_cpu): boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance']) boxes = boxes * scale / resize boxes = boxes.cpu().numpy() - scores = conf.squeeze(0).data.cpu().numpy()[:, 1] + scores = conf.squeeze(0).data.cpu().numpy() landms = decode_landm(landms.data.squeeze(0), prior_data, cfg['variance']) scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2],