def create_mtcnn_net(image,
mini_face,
device,
p_model_path=None,
r_model_path=None,
o_model_path=None):
boxes = np.array([])
landmarks = np.array([])
if p_model_path is not None:
pnet = PNet().to(device)
pnet.load_state_dict(
torch.load(p_model_path,
map_location=lambda storage, loc: storage))
pnet.eval()
bboxes = detect_pnet(pnet, image, mini_face, device)
if r_model_path is not None:
rnet = RNet().to(device)
rnet.load_state_dict(
torch.load(r_model_path,
map_location=lambda storage, loc: storage))
rnet.eval()
bboxes = detect_rnet(rnet, image, bboxes, device)
if o_model_path is not None:
onet = ONet().to(device)
onet.load_state_dict(
torch.load(o_model_path,
map_location=lambda storage, loc: storage))
onet.eval()
bboxes, landmarks = detect_onet(onet, image, bboxes, device)
return bboxes, landmarks
batch_size=batch_size,
shuffle=True),
'val':
torch.utils.data.DataLoader(ListDataset(val_path),
batch_size=batch_size,
shuffle=True)
}
dataset_sizes = {
'train': len(ListDataset(train_path)),
'val': len(ListDataset(val_path))
}
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# load the model and weights for initialization
model = ONet(is_train=True).to(device)
model.apply(weights_init)
optimizer = torch.optim.Adam(model.parameters())
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_accuracy = 0.0
best_loss = 100
loss_cls = nn.CrossEntropyLoss()
loss_offset = nn.MSELoss()
loss_landmark = nn.MSELoss()
num_epochs = 16
for epoch in range(num_epochs):
__all__ = ['TestEngineImg']
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
p_model_path = 'models/pnet_Weights'
o_model_path = 'models/onet_Weights'
l_model_path = 'models/landmark.pkl'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
pnet = PNet().to(device)
pnet.load_state_dict(
torch.load(p_model_path, map_location=lambda storage, loc: storage))
pnet.eval()
onet = ONet().to(device)
onet.load_state_dict(
torch.load(o_model_path, map_location=lambda storage, loc: storage))
onet.eval()
landmark = torch.load(l_model_path)
landmark = landmark.to(device)
landmark.eval()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='MTCNN Demo')
parser.add_argument("--test_image",
dest='test_image',
help="test image path",
default="./lmark",
del model.conv4_1
del model.conv4_2
model.conv4_1 = new_conv4_1
model.conv4_2 = new_conv4_2
return model
if __name__ == '__main__':
import sys
sys.path.append("../Base_Model")
from MTCNN_nets import PNet, RNet, ONet
model = ONet()
model.train()
layer_index = 9
filter_index = (2,4)
model = prune_mtcnn(model, layer_index, *filter_index, use_cuda=False)
print(model)