img2pose: Face Alignment and Detection via 6DoF, Face Pose Estimation代码理解

import argparse
import os
import sys
import time

import numpy as np
from PIL import Image, ImageOps
from torchvision import transforms
from tqdm import tqdm
import cv2

sys.path.append("./")
from img2pose import img2poseModel
from model_loader import load_model


def init_model(pretrained_path, threed_68_points, pose_stddev, pose_mean, depth):
    img2pose_model = img2poseModel(
            depth,
            200,
            1400,
            pose_mean=np.load(pose_mean),
            pose_stddev=np.load(pose_stddev),
            threed_68_points=np.load(threed_68_points),
        )
    load_model(
            img2pose_model.fpn_model,
            pretrained_path,
            cpu_mode=str(img2pose_model.device) == "cpu",
            model_only=True,
        )
    img2pose_model.evaluate()

    return img2pose_model


def bbox_voting(bboxes, iou_thresh=0.6):
    # bboxes: a numpy array of N*5 size representing N boxes;
    #         for each box, it is represented as [x1, y1, x2, y2, s]
    # iou_thresh: group bounding boxes if their overlap is > threshold.
    bboxes = np.asarray(bboxes)
    order = bboxes[:, 4].ravel().argsort()[::-1]
    bboxes = bboxes[order, :]
    areas = (bboxes[:, 2] - bboxes[:, 0] + 1) * (bboxes[:, 3] - bboxes[:, 1] + 1)
    voted_bboxes = np.zeros([0, 5])
    while bboxes.shape[0] > 0:
        xx1 = np.maximum(bboxes[0, 0], bboxes[:, 0])
        yy1 = np.maximum(bboxes[0, 1], bboxes[:, 1])
        xx2 = np.minimum(bboxes[0, 2], bboxes[:, 2])
        yy2 = np.minimum(bboxes[0, 3], b