facexlib

1
import argparse
2
import cv2
3
import torch
4

5
from facexlib.detection import init_detection_model
6
from facexlib.visualization import visualize_detection
7

8

9
def main(args):
10
    # initialize model
11
    det_net = init_detection_model(args.model_name, half=args.half)
12

13
    img = cv2.imread(args.img_path)
14
    with torch.no_grad():
15
        bboxes = det_net.detect_faces(img, 0.97)
16
        # x0, y0, x1, y1, confidence_score, five points (x, y)
17
        print(bboxes)
18
        visualize_detection(img, bboxes, args.save_path)
19

20

21
if __name__ == '__main__':
22
    parser = argparse.ArgumentParser()
23
    parser.add_argument('--img_path', type=str, default='assets/test.jpg')
24
    parser.add_argument('--save_path', type=str, default='test_detection.png')
25
    parser.add_argument(
26
        '--model_name', type=str, default='retinaface_resnet50', help='retinaface_resnet50 | retinaface_mobile0.25')
27
    parser.add_argument('--half', action='store_true')
28
    args = parser.parse_args()
29

30
    main(args)
31