File size: 2,104 Bytes
f12ab4c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import face_alignment
import cv2
import os
from os.path import join
import numpy as np
from tqdm import tqdm
import json
from glob import glob
import argparse
class Gen2DLandmarks(object):
def __init__(self) -> None:
super().__init__()
#
# print(face_alignment)
# print(face_alignment.FaceAlignment)
# print(face_alignment.LandmarksType)
self.fa_func = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)
# self.fa_func = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False)
def main_process(self, img_dir):
img_path_list = [x for x in sorted(glob("%s/*.png" % img_dir)) if "mask" not in x]
img_path_list = img_path_list + [x for x in sorted(glob("%s/*.jpg" % img_dir)) if "mask" not in x]
#img_path_list = img_path_list[27:]
if len(img_path_list) == 0:
print("Dir: %s does include any .png and .jpg images." % img_dir)
exit(0)
img_path_list.sort()
for img_path in tqdm(img_path_list, desc="Generate facial landmarks"):
img_bgr = cv2.imread(img_path)
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
res = self.fa_func.get_landmarks(img_rgb)
if res is None:
print("Warning: can't predict the landmark info of %s" % img_path)
# base_name = img_path[img_path.rfind("/") + 1:-4]
save_path = img_path[:-4] + "_lm2d.txt"
preds = res[0]
with open(save_path, "w") as f:
for tt in preds:
f.write("%f %f\n"%(tt[0],tt[1]))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='The code for generating facial landmarks.')
# parser.add_argument("--gpu_id", type=int, default=0)
parser.add_argument("--img_dir", type=str, required=True)
args = parser.parse_args()
tt = Gen2DLandmarks()
tt.main_process(args.img_dir)
|