facexlib

1
import argparse
2
import cv2
3
import torch
4

5
from facexlib.alignment import init_alignment_model, landmark_98_to_68
6
from facexlib.visualization import visualize_alignment
7

8

9
def main(args):
10
    # initialize model
11
    align_net = init_alignment_model(args.model_name, device=args.device)
12

13
    img = cv2.imread(args.img_path)
14
    with torch.no_grad():
15
        landmarks = align_net.get_landmarks(img)
16
        if args.to68:
17
            landmarks = landmark_98_to_68(landmarks)
18
        visualize_alignment(img, [landmarks], args.save_path)
19

20

21
if __name__ == '__main__':
22
    parser = argparse.ArgumentParser()
23
    parser.add_argument('--img_path', type=str, default='assets/test2.jpg')
24
    parser.add_argument('--save_path', type=str, default='test_alignment.png')
25
    parser.add_argument('--model_name', type=str, default='awing_fan')
26
    parser.add_argument('--device', type=str, default='cuda')
27
    parser.add_argument('--to68', action='store_true')
28
    args = parser.parse_args()
29

30
    main(args)
31