def test_onet(self):
onet = mtcnn.ONet().to_script()
data = torch.randn(100, 3, 48, 48)
det, box, landmarks = onet(data)
self.assertEqual(list(det.shape), [100, 2])
self.assertEqual(list(box.shape), [100, 4])
self.assertEqual(list(landmarks.shape), [100, 10])
def get_net():
"""
Create pnet, rnet, onet for detector.
"""
pnet = mtcnn_pytorch.PNet()
rnet = mtcnn_pytorch.RNet()
onet = mtcnn_pytorch.ONet()
return pnet, rnet, onet
def test_load_caffe_model(self):
pnet = mtcnn.PNet()
rnet = mtcnn.RNet()
onet = mtcnn.ONet()
weight_folder = os.path.join(here, '../output/converted')
pnet.load_caffe_model(
np.load(os.path.join(weight_folder, 'pnet.npy'))[()])
rnet.load_caffe_model(
np.load(os.path.join(weight_folder, 'rnet.npy'))[()])
onet.load_caffe_model(
np.load(os.path.join(weight_folder, 'onet.npy'))[()])
def test_onet(self):
onet = mtcnn.ONet(is_train=True)
data = torch.randn(100, 3, 48, 48)
det, box, landmarks = onet(data)
self.assertEqual(list(det.shape), [100, 2])
self.assertEqual(list(box.shape), [100, 4])
self.assertEqual(list(landmarks.shape), [100, 10])
onet.get_loss(data, torch.ones(100, dtype=torch.int64),
torch.randn(100, 4), torch.randn(100, 10))
def get_net(weight_folder=None):
"""
Create pnet, rnet, onet for detector.
"""
pnet = mtcnn_pytorch.PNet()
rnet = mtcnn_pytorch.RNet()
onet = mtcnn_pytorch.ONet()
if weight_folder is not None:
pnet.load(os.path.join(weight_folder, 'pnet'))
rnet.load(os.path.join(weight_folder, 'rnet'))
onet.load(os.path.join(weight_folder, 'onet'))
return pnet, rnet, onet
def setUp(self):
weight_folder = os.path.join(here, '../output/converted')
pnet = mtcnn.PNet()
rnet = mtcnn.RNet()
onet = mtcnn.ONet()
pnet.load_caffe_model(
np.load(os.path.join(weight_folder, 'pnet.npy'))[()])
rnet.load_caffe_model(
np.load(os.path.join(weight_folder, 'rnet.npy'))[()])
onet.load_caffe_model(
np.load(os.path.join(weight_folder, 'onet.npy'))[()])
self.detector = detect.FaceDetector(pnet, rnet, onet, "cuda:0")
self.test_img = os.path.join(here, 'asset/images/office5.jpg')
def setUp(self):
self.test_video = os.path.join(here, './asset/video/school.avi')
weight_folder = os.path.join(here, '../output/converted')
pnet = mtcnn.PNet()
rnet = mtcnn.RNet()
onet = mtcnn.ONet()
pnet.load_caffe_model(
np.load(os.path.join(weight_folder, 'pnet.npy'))[()])
rnet.load_caffe_model(
np.load(os.path.join(weight_folder, 'rnet.npy'))[()])
onet.load_caffe_model(
np.load(os.path.join(weight_folder, 'onet.npy'))[()])
self.detector = FaceDetector(pnet, rnet, onet)
self.tracker = FaceTracker(self.detector)