|
|
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 |
