|
- import cv2
- from core.detect import create_mtcnn_net, MtcnnDetector
- from core.vision import vis_face
-
- if __name__ == '__main__':
- # original model
-
- # p_model_path = "./original_model/pnet_epoch.pt"
- # r_model_path = "./original_model/rnet_epoch.pt"
- # o_model_path = "./original_model/onet_epoch.pt "
-
- p_model_path = "./model_store/pnet_epoch.pt"
- r_model_path = "./model_store/rnet_epoch.pt"
- o_model_path = "./model_store/onet_epoch.pt"
-
- pnet, rnet, onet = create_mtcnn_net(p_model_path=p_model_path, r_model_path=r_model_path, o_model_path=o_model_path,
- use_cuda=False)
- mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=onet, min_face_size=12, threshold=[0.96, 0.03, 0.005])
-
- img = cv2.imread("img_for_test/mid_new.jpg")
- img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- bboxs, landmarks = mtcnn_detector.detect_face(img)
- bbox_num = bboxs.shape[0]
- print("bboxs : ", bbox_num)
- save_name = './test_result/' + str(bbox_num) + '_r.jpg'
- vis_face(img_bg, bboxs, landmarks, save_name)
-
- # #original model
- # o_model_path = "./original_model/onet_epoch.pt"
-
- # #trained model
- # p_model_path = "./model_store/pnet_epoch_11.pt"
- # r_model_path = "./model_store/rnet_epoch_9.pt"
-
- # pnet, rnet, onet = create_mtcnn_net(p_model_path=p_model_path, r_model_path=r_model_path, o_model_path=o_model_path, use_cuda=False)
- # mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=onet, min_face_size=12, threshold=[0.3, 0.1, 0.3])
-
- # img = cv2.imread("mid.jpg")
- # img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- # bboxs, landmarks = mtcnn_detector.detect_face(img)
-
- # save_name = 'r_1.jpg'
- # vis_face(img_bg,bboxs,landmarks, save_name)
|