yolov8-face 带五个关键点的推理代码

学习 使用 GitHub - derronqi/yolov8-face: yolov8 face detection with landmark的推理代码

import cv2
import numpy as np
import math
import argparse
import time


class YOLOv8_face:
    def __init__(self, path, conf_thres=0.25, iou_thres=0.45):
        self.conf_threshold = conf_thres
        self.iou_threshold = iou_thres
        self.class_names = ['face']
        self.num_classes = len(self.class_names)
        # Initialize model
        # self.net = cv2.dnn.readNet(path)
        self.net = cv2.dnn.readNetFromONNX(path)
        self.input_height = 640
        self.input_width = 640

        self.strides = (8, 16, 32)
        self.feats_hw = [(math.ceil(self.input_height / self.strides[i]), math.ceil(self.input_width / self.strides[i]))
                         for i in range(len(self.strides))]

    def make_anchors(self, feats_hw, grid_cell_offset=0.5):
        """Generate anchors from features."""
        anchor_points = {}
        for i, stride in enumerate(self.strides):
            h, w = feats_hw[i]
            x = np.arange(0, w) + grid_cell_offset  # shift x
            y = np.arange(0, h) + grid_cell_offset  # shift y
            sx, sy = np.meshgrid(x, y)
            # sy, sx = np.meshgrid(y, x)
            anchor_points[stride] = np.stack((sx, sy), axis=-1).reshape(-1, 2)
        return anchor_points

    def softmax(self, x, axis=1):
        x_exp = np.exp(x)
        # 如果是列向量，则axis=0
        x_sum = np.sum(x_exp, axis=axis, keepdims=True)
        s = x_exp / x_sum
        return s

    def resize_image(self, srcimg, keep_ratio=True):
        top, left, newh, neww = 0, 0, self.input_width, self.input_height
        if keep_ratio and srcimg.shape[0] != srcimg.shape[1]:
            hw_scale = srcimg.shape[0] / srcimg.shape[1]
            if hw_scale > 1:
                newh, neww = self.input_height, int(self.input_width / hw_scale)
                img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA)
                left = int((self.input_width - neww) * 0.5)
                img = cv2.copyMakeBorder(img, 0, 0, left, self.input_width - neww - left, cv2.BORDER_CONSTANT,
                                         value=(0, 0, 0))  # add border
            else:
                newh, neww = int(self.input_height * hw_scale), self.input_width
                img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA)
                top = int((self.input_height - newh) * 0.5)
                img = cv2.copyMakeBorder(img, top, self.input_height - newh - top, 0, 0, cv2.BORDER_CONSTANT,
                                         value=(0, 0, 0))
        else:
            img = cv2.resize(srcimg, (self.input_width, self.input_height), interpolation=cv2.INTER_AREA)
        return img, newh, neww, top, left

    def detect(self, srcimg):
        [height, width, _] = srcimg.shape
        length = max((height, width))
        image = np.zeros((length, length, 3), np.uint8)
        image[0:height, 0:width] = srcimg
        scale = length / 640

        blob = cv2.dnn.blobFromImage(image, scalefactor=1 / 255, size=(640, 640), swapRB=True)
        self.net.setInput(blob)
        preds = self.net.forward()
        # if isinstance(preds, tuple):
        #     preds = list(preds)
        # if float(cv2.__version__[:3])>=4.7:
        #     preds = [preds[2], preds[0], preds[1]] ###opencv4.7需要这一步，opencv4.5不需要
        # Perform inference on the image
        det_bboxes, det_conf, landmarks = self.post_process(preds, scale)
        return det_bboxes, det_conf, landmarks

    def post_process_1(self, preds, scale):
        preds = preds.transpose((0, 2, 1)).squeeze()

        t_index = np.where(preds[:, 4] > 0.25)
        preds = preds[t_index, :].squeeze()

        x1 = np.expand_dims(preds[:, 0] - (0.5 * preds[:, 2]), -1)
        y1 = np.expand_dims(preds[:, 1] - (0.5 * preds[:, 3]), -1)
        w = np.expand_dims(preds[:, 2], -1)
        h = np.expand_dims(preds[:, 3], -1)
        bboxes = np.concatenate([x1, y1, w, h], axis=-1) * scale
        scores = preds[:, 4]
        landmarks = preds[:, -15:] * scale
        indices = cv2.dnn.NMSBoxes(bboxes, scores, 0.25, 0.45, 0.5)

        if len(indices) > 0:
            bboxes = bboxes[indices]
            scores = scores[indices]
            # classIds = classIds[indices]
            landmarks = landmarks[indices]
            return bboxes, scores, landmarks
        else:
            print('nothing detect')

    def post_process(self, preds, scale):
        bboxes, scores, landmarks = [], [], []
        preds = preds.transpose((0, 2, 1))

        rows = preds.shape[1]
        for i in range(rows):
            score = preds[0][i][4]
            kpt = preds[0][i][-15:]
            if score >= 0.25:
                # xywh
                box = [
                    preds[0][i][0] - (0.5 * preds[0][i][2]), preds[0][i][1] - (0.5 * preds[0][i][3]),
                    preds[0][i][2], preds[0][i][3]]
                bboxes.append(box)
                scores.append(score)
                landmarks.append(kpt)

        bboxes = np.array(bboxes) * scale
        scores = np.array(scores)
        landmarks = np.array(landmarks) * scale
        indices = cv2.dnn.NMSBoxes(bboxes, scores, 0.25, 0.45, 0.5)

        if len(indices) > 0:
            bboxes = bboxes[indices]
            scores = scores[indices]
            # classIds = classIds[indices]
            landmarks = landmarks[indices]
            return bboxes, scores, landmarks
        else:
            print('nothing detect')

    def draw_detections(self, image, boxes, scores, kpts):
        for box, score, kp in zip(boxes, scores, kpts):
            x, y, w, h = box.astype(int)
            # Draw rectangle
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), thickness=2)
            cv2.putText(image, "face:" + str(round(score, 2)), (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255),
                        thickness=2)
            for i in range(5):
                cv2.circle(image, (int(kp[i * 3]), int(kp[i * 3 + 1])), 1, (0, 255, 0), thickness=-1)
                # cv2.putText(image, str(i), (int(kp[i * 3]), int(kp[i * 3 + 1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), thickness=1)
        return image


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgpath', type=str, default='images/jihui.png', help="image path")  # big_face.png
    parser.add_argument('--modelpath', type=str, default='weights/yolov8n-face.onnx',
                        help="onnx filepath")
    parser.add_argument('--confThreshold', default=0.45, type=float, help='class confidence')
    parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh')
    args = parser.parse_args()

    # Initialize YOLOv8_face object detector
    YOLOv8_face_detector = YOLOv8_face(args.modelpath, conf_thres=args.confThreshold, iou_thres=args.nmsThreshold)
    srcimg: np.ndarray = cv2.imread(args.imgpath)

    # Detect Objects
    boxes, scores, kpts = YOLOv8_face_detector.detect(srcimg)

    # Draw detections
    dstimg = YOLOv8_face_detector.draw_detections(srcimg, boxes, scores, kpts)
    cv2.imwrite('result.jpg', dstimg)