| import os |
| import sys |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.append(BASE_DIR) |
| import torch |
| from src.dataset.face_align.yoloface import YoloFace |
|
|
| class AlignImage(object): |
| def __init__(self, device='cuda', det_path='checkpoints/yoloface_v5m.pt'): |
| self.facedet = YoloFace(pt_path=det_path, confThreshold=0.5, nmsThreshold=0.45, device=device) |
|
|
| @torch.no_grad() |
| def __call__(self, im, maxface=False): |
| bboxes, kpss, scores = self.facedet.detect(im) |
| face_num = bboxes.shape[0] |
|
|
| five_pts_list = [] |
| scores_list = [] |
| bboxes_list = [] |
| for i in range(face_num): |
| five_pts_list.append(kpss[i].reshape(5,2)) |
| scores_list.append(scores[i]) |
| bboxes_list.append(bboxes[i]) |
|
|
| if maxface and face_num>1: |
| max_idx = 0 |
| max_area = (bboxes[0, 2])*(bboxes[0, 3]) |
| for i in range(1, face_num): |
| area = (bboxes[i,2])*(bboxes[i,3]) |
| if area>max_area: |
| max_idx = i |
| five_pts_list = [five_pts_list[max_idx]] |
| scores_list = [scores_list[max_idx]] |
| bboxes_list = [bboxes_list[max_idx]] |
|
|
| return five_pts_list, scores_list, bboxes_list |
