Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
hackassin
GitHub Repository: hackassin/learnopencv
 Path: blob/master/Face-Recognition-with-ArcFace/face_alignment.py
3118 views
1
# Original code

2
# https://github.com/ZhaoJ9014/face.evoLVe.PyTorch/blob/master/align/face_align.py

3


4
import argparse

5
import os

6


7
import numpy as np

8
from align.align_trans import (

9
    get_reference_facial_points,

10
    warp_and_crop_face,

11
)

12
from align.detector import detect_faces

13
from PIL import Image

14
from tqdm import tqdm

15


16
if __name__ == "__main__":

17
    parser = argparse.ArgumentParser()

18
    parser.add_argument(

19
        "--tags",

20
        help="specify your tags for raw datasets",

21
        default="test",

22
        nargs='+',

23
        required=True

24
    )

25
    parser.add_argument(

26
        "--crop_size",

27
        help="specify size of aligned faces",

28
        default=112,

29
        choices=[112, 224],

30
        type=int,

31
    )

32
    args = parser.parse_args()

33


34
    tags = args.tags

35
    crop_size = args.crop_size

36
    scale = crop_size / 112.0

37
    reference = get_reference_facial_points(default_square=True) * scale

38


39
    for tag in tags:

40
        source_root = os.path.join("data", tag)

41
        dest_root = source_root + "_aligned"

42
        if not os.path.isdir(dest_root):

43
            os.mkdir(dest_root)

44


45
        for subfolder in tqdm(os.listdir(source_root)):

46
            if not os.path.isdir(os.path.join(dest_root, subfolder)):

47
                os.mkdir(os.path.join(dest_root, subfolder))

48
            for image_name in os.listdir(os.path.join(source_root, subfolder)):

49
                print(

50
                    "Processing\t{}".format(

51
                        os.path.join(source_root, subfolder, image_name),

52
                    ),

53
                )

54
                img = Image.open(os.path.join(source_root, subfolder, image_name))

55
                try:  # Handle exception

56
                    _, landmarks = detect_faces(img)

57
                except Exception:

58
                    print(

59
                        "{} is discarded due to exception!".format(

60
                            os.path.join(source_root, subfolder, image_name),

61
                        ),

62
                    )

63
                    continue

64
                if (

65
                    len(landmarks) == 0

66
                ):  # If the landmarks cannot be detected, the img will be discarded

67
                    print(

68
                        "{} is discarded due to non-detected landmarks!".format(

69
                            os.path.join(source_root, subfolder, image_name),

70
                        ),

71
                    )

72
                    continue

73
                facial5points = [[landmarks[0][j], landmarks[0][j + 5]] for j in range(5)]

74
                warped_face = warp_and_crop_face(

75
                    np.array(img),

76
                    facial5points,

77
                    reference,

78
                    crop_size=(crop_size, crop_size),

79
                )

80
                img_warped = Image.fromarray(warped_face)

81
                if image_name.split(".")[-1].lower() not in ["jpg", "jpeg"]:

82
                    image_name = ".".join(image_name.split(".")[:-1]) + ".jpg"

83
                img_warped.save(os.path.join(dest_root, subfolder, image_name))

84


85