head_detect/demo.py

# -*- coding:utf-8 –*-
import os
import sys

import cv2

sys.path.append(os.path.dirname(os.path.dirname(sys.path[0])))
from head_detect.head_detector.head_detectorv4 import HeadDetector
from head_detect.utils_quailty_assurance.utils_quailty_assurance import find_files, write_json
from utils.images import crop_face_square_rate
import numpy as np

ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
model_path = f"{ROOT_DIR}/models/HeadDetectorv1.6.onnx"

headDetector = HeadDetector(onnx_path=model_path)


def detect_driver_face(img, show_res=False, show_crop=False):
    if isinstance(img, str):
        img = cv2.imread(img)
    else:
        img = img.copy()
    image_heigth,image_width = img.shape[:2]
    # origin_image = cv2.resize(img, (1280, 720))
    # origin_image = img[:, w//2:, :]
    # short_side = min(image_heigth,image_width)
    width_shift,heigth_shift = image_width//2,0
    # cv2.imwrite("squre.jpg",img[heigth_shift:,width_shift:,:])
    bboxes = headDetector.run(img[:,width_shift:, :], get_largest=True)  # 人脸检测，获取面积最大的人脸
    if not bboxes:
        return [0, 0, 0, 0], 0
    box = [int(width_shift+bboxes[0][0]),int(heigth_shift+bboxes[0][1]),int(width_shift+bboxes[0][2]),int(heigth_shift+bboxes[0][3])]

    # box[0], box[1] = max(0, box[0]), max(0, box[1])
    # box[2], box[3] = min(w, box[2]), min(h, box[3])
    score = bboxes[0][-1]

    if (box[2] - box[0]) == 0 or (box[3] - box[1]) == 0:
        return [0, 0, 0, 0], 0

    # print(box, pred, score)
    if show_res:
        x0, y0, x1, y1 = box
        print(box)
        cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), thickness=2)
        if show_crop:
            _, crop_box = crop_face_square_rate(img, box, rate=-0.07)
            cx0, cy0, cx1, cy1 = crop_box
            print(crop_box)
            cv2.rectangle(img, (cx0, cy0), (cx1, cy1), (0, 255, 0), thickness=2)
        cv2.imshow('res', img)
        cv2.waitKey(0)

    return box, score


if __name__ == "__main__":
    detect_driver_face('../../test_visual/look_down.jpg', show_res=True, show_crop=True)