检测效果图

可支持网络图片 base64 本地路径方式
代码可以通过fastapi 改成服务启动模式

实现代码 可支持网络图片 base64 本地路径方式

import json
import platform
import time
import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import cv2
import base64
import numpy as np
import requests
from scrfd import SCRFD
import onnxruntime as ort
from urllib.parse import urlparse

ort.set_default_logger_severity(3)  # 忽略警告

# 获取模型文件所在目录的路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
onnx_path = os.path.join(BASE_DIR, "..", "model.onnx")
onnx_path = os.path.abspath(onnx_path)

# 用于存储 SCRFD 检测器
detector = None
detector = SCRFD(onnx_path)
detector.prepare(1)  # 开启GPU加速
print("====加载模型====")

def is_web_url(path):
    """判断是否为网络URL"""
    try:
        result = urlparse(path)
        return all([result.scheme in ('http', 'https'), result.netloc])
    except:
        return False


# 读取图片流
def read_image(source):
    """通用图片读取函数"""
    start_time = time.time()
    try:
        if is_web_url(source):
            # 网络图片处理
            response = requests.get(source, timeout=10, stream=True)
            response.raise_for_status()
            img_array = np.frombuffer(response.content, np.uint8)
            image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        else:
            # 本地图片处理
            if not os.path.exists(source):
                raise FileNotFoundError(f"本地文件不存在: {source}")
                # 加载图片
            if platform.system() == "Windows":
                image = cv2.imdecode(np.fromfile(source, dtype=np.uint8), cv2.IMREAD_COLOR)
            else:
                image = cv2.imread(source)
            if image is None:
                raise ValueError(f"无法读取图片文件: {source}")

        return image
    except Exception as e:
        print(f"图片加载失败: {source}, 错误: {e}")
        return None
    finally:
        print(f"图片加载耗时: {time.time() - start_time:.3f} 秒")


def process_images_from_txt(image_url):
    image = read_image(image_url)  # image_url是http/https链接
    if image is None:
        print(f"无法读取图片: {image_url}")
        return 'error'

    start_time = time.time()  # 计算开始时间
    bboxes, kpss = detector.autodetect(image, 32)

    # 将检测结果转换为字典格式
    face_data = [
        {
            'topLeftY': round(float(box[1]), 2),
            'topLeftX': round(float(box[0]), 2),
            'lowerRightY': round(float(box[3]), 2),
            'lowerRightX': round(float(box[2]), 2),
            'confidenceLevel': round(float(box[4]), 2)
        }
        for box in bboxes
    ]

    print(f"检测耗时:{image_url} {time.time() - start_time:.3f} 秒")
    print(face_data)
    # # 可视化人脸框
    for box in bboxes:
        x1, y1, x2, y2, score = box
        color = (0, 255, 0)  # 绿色
        thickness = 2
        cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
        cv2.putText(image, f"{score:.2f}", (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    # 保存可视化图片（可选）
    output_path = os.path.join(BASE_DIR, f"{time.time()}.jpg")
    cv2.imwrite(output_path, image)
    print(f"已保存结果图像到: {output_path}")
    return json.dumps(face_data)


def process_images_from_base64(base_str):
    """
    从 Base64 字符串读取图片并进行人脸检测
    :param base_str: 图片的 Base64 编码字符串（不含前缀如 "data:image/jpeg;base64,"）
    :return: 人脸检测结果的 JSON 字符串
    """
    try:
        # 解码 Base64 字符串为二进制数据
        img_data = base64.b64decode(base_str)

        # 将二进制数据转换为 numpy 数组（OpenCV 格式）
        img_array = np.frombuffer(img_data, dtype=np.uint8)

        # 用 OpenCV 读取为图像（BGR 格式）
        image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        if image is None:
            return json.dumps({"error": "无法解析 Base64 图片数据"})

        # 人脸检测（保持原逻辑）
        start_time = time.time()
        bboxes, kpss = detector.autodetect(image, 32)  # 假设 detector 是已初始化的检测模型
        print(f"检测耗时: {time.time() - start_time:.3f} 秒")

        # 格式化检测结果
        face_data = [
            {
                'topLeftY': round(float(box[1]), 2),
                'topLeftX': round(float(box[0]), 2),
                'lowerRightY': round(float(box[3]), 2),
                'lowerRightX': round(float(box[2]), 2),
                'confidenceLevel': round(float(box[4]), 2)
            }
            for box in bboxes
        ]
        print("检测结果:", face_data)

        # 可视化人脸框
        for box in bboxes:
            x1, y1, x2, y2, score = box
            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            cv2.putText(image, f"{score:.2f}", (int(x1), int(y1) - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        # # 保存可视化图片（可选）
        output_path = os.path.join(BASE_DIR, f"{time.time()}.jpg")
        cv2.imwrite(output_path, image)
        print(f"已保存结果图像到: {output_path}")
        # 返回检测结果 JSON
        return json.dumps(face_data)

    except Exception as e:
        print(e)
        print(f"处理失败: {str(e)}")
        return json.dumps({"error": str(e)})


def image_to_base64(image_path):
    """将本地图片转换为 Base64 字符串（不含前缀）"""
    with open(image_path, "rb") as f:
        # 读取图片二进制数据并编码为Base64
        base64_str = base64.b64encode(f.read()).decode("utf-8")
    return base64_str


if __name__ == "__main__":
    url = '目标图片路径.jpg'
    # url = 'http://目标图片链接.jpg'
    faceRecognitionJson = process_images_from_txt(url)
    print(faceRecognitionJson)

    # base64str = image_to_base64('/Users/haoruijie/Downloads/幼儿园1班级照/BL6Z0170.jpg')
    # base64检测
    # faceRecognitionJson = process_images_from_base64(base64str)

scrfd文件代码

from __future__ import division
import numpy as np
import os.path as osp
import cv2

def softmax(z):
    assert len(z.shape) == 2
    s = np.max(z, axis=1)
    s = s[:, np.newaxis] # necessary step to do broadcasting
    e_x = np.exp(z - s)
    div = np.sum(e_x, axis=1)
    div = div[:, np.newaxis] # dito
    return e_x / div

def distance2bbox(points, distance, max_shape=None):
    """Decode distance prediction to bounding box.

    Args:
        points (Tensor): Shape (n, 2), [x, y].
        distance (Tensor): Distance from the given point to 4
            boundaries (left, top, right, bottom).
        max_shape (tuple): Shape of the image.

    Returns:
        Tensor: Decoded bboxes.
    """
    x1 = points[:, 0] - distance[:, 0]
    y1 = points[:, 1] - distance[:, 1]
    x2 = points[:, 0] + distance[:, 2]
    y2 = points[:, 1] + distance[:, 3]
    if max_shape is not None:
        x1 = x1.clamp(min=0, max=max_shape[1])
        y1 = y1.clamp(min=0, max=max_shape[0])
        x2 = x2.clamp(min=0, max=max_shape[1])
        y2 = y2.clamp(min=0, max=max_shape[0])
    return np.stack([x1, y1, x2, y2], axis=-1)

def distance2kps(points, distance, max_shape=None):
    """Decode distance prediction to bounding box.

    Args:
        points (Tensor): Shape (n, 2), [x, y].
        distance (Tensor): Distance from the given point to 4
            boundaries (left, top, right, bottom).
        max_shape (tuple): Shape of the image.

    Returns:
        Tensor: Decoded bboxes.
    """
    preds = []
    for i in range(0, distance.shape[1], 2):
        px = points[:, i%2] + distance[:, i]
        py = points[:, i%2+1] + distance[:, i+1]
        if max_shape is not None:
            px = px.clamp(min=0, max=max_shape[1])
            py = py.clamp(min=0, max=max_shape[0])
        preds.append(px)
        preds.append(py)
    return np.stack(preds, axis=-1)

class SCRFD:
    def __init__(self, model_file=None, session=None):
        import onnxruntime
        self.model_file = model_file
        self.session = session
        self.taskname = 'detection'
        self.batched = False
        if self.session is None:
            assert self.model_file is not None
            assert osp.exists(self.model_file)
            self.session = onnxruntime.InferenceSession(self.model_file, providers=['CUDAExecutionProvider'])
        self.center_cache = {}
        self.nms_thresh = 0.4
        self.det_thresh = 0.5
        self._init_vars()

    def _init_vars(self):
        input_cfg = self.session.get_inputs()[0]
        input_shape = input_cfg.shape
        #print(input_shape)
        if isinstance(input_shape[2], str):
            self.input_size = None
        else:
            self.input_size = tuple(input_shape[2:4][::-1])
        #print('image_size:', self.image_size)
        input_name = input_cfg.name
        self.input_shape = input_shape
        outputs = self.session.get_outputs()
        if len(outputs[0].shape) == 3:
            self.batched = True
        output_names = []
        for o in outputs:
            output_names.append(o.name)
        self.input_name = input_name
        self.output_names = output_names
        self.input_mean = 127.5
        self.input_std = 128.0
        #print(self.output_names)
        #assert len(outputs)==10 or len(outputs)==15
        self.use_kps = False
        self._anchor_ratio = 1.0
        self._num_anchors = 1
        if len(outputs)==6:
            self.fmc = 3
            self._feat_stride_fpn = [8, 16, 32]
            self._num_anchors = 2
        elif len(outputs)==9:
            self.fmc = 3
            self._feat_stride_fpn = [8, 16, 32]
            self._num_anchors = 2
            self.use_kps = True
        elif len(outputs)==10:
            self.fmc = 5
            self._feat_stride_fpn = [8, 16, 32, 64, 128]
            self._num_anchors = 1
        elif len(outputs)==15:
            self.fmc = 5
            self._feat_stride_fpn = [8, 16, 32, 64, 128]
            self._num_anchors = 1
            self.use_kps = True

    def prepare(self, ctx_id, **kwargs):
        if ctx_id<0:
            self.session.set_providers(['CPUExecutionProvider'])
        if ctx_id == 1: #gpu
            self.session.set_providers(['CUDAExecutionProvider'])
        nms_thresh = kwargs.get('nms_thresh', None)
        if nms_thresh is not None:
            self.nms_thresh = nms_thresh
        det_thresh = kwargs.get('det_thresh', None)
        if det_thresh is not None:
            self.det_thresh = det_thresh
        input_size = kwargs.get('input_size', None)
        if input_size is not None:
            if self.input_size is not None:
                print('warning: det_size is already set in scrfd model, ignore')
            else:
                self.input_size = input_size

    def forward(self, img, threshold):
        scores_list = []
        bboxes_list = []
        kpss_list = []
        input_size = tuple(img.shape[0:2][::-1])
        blob = cv2.dnn.blobFromImage(img, 1.0/self.input_std, input_size, (self.input_mean, self.input_mean, self.input_mean), swapRB=True)
        net_outs = self.session.run(self.output_names, {self.input_name : blob})

        input_height = blob.shape[2]
        input_width = blob.shape[3]
        fmc = self.fmc
        for idx, stride in enumerate(self._feat_stride_fpn):
            # If model support batch dim, take first output
            if self.batched:
                scores = net_outs[idx][0]
                bbox_preds = net_outs[idx + fmc][0]
                bbox_preds = bbox_preds * stride
                if self.use_kps:
                    kps_preds = net_outs[idx + fmc * 2][0] * stride
            # If model doesn't support batching take output as is
            else:
                scores = net_outs[idx]
                bbox_preds = net_outs[idx + fmc]
                bbox_preds = bbox_preds * stride
                if self.use_kps:
                    kps_preds = net_outs[idx + fmc * 2] * stride

            height = input_height // stride
            width = input_width // stride
            K = height * width
            key = (height, width, stride)
            if key in self.center_cache:
                anchor_centers = self.center_cache[key]
            else:
                #solution-1, c style:
                #anchor_centers = np.zeros( (height, width, 2), dtype=np.float32 )
                #for i in range(height):
                #    anchor_centers[i, :, 1] = i
                #for i in range(width):
                #    anchor_centers[:, i, 0] = i

                #solution-2:
                #ax = np.arange(width, dtype=np.float32)
                #ay = np.arange(height, dtype=np.float32)
                #xv, yv = np.meshgrid(np.arange(width), np.arange(height))
                #anchor_centers = np.stack([xv, yv], axis=-1).astype(np.float32)

                #solution-3:
                anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32)
                #print(anchor_centers.shape)

                anchor_centers = (anchor_centers * stride).reshape( (-1, 2) )
                if self._num_anchors>1:
                    anchor_centers = np.stack([anchor_centers]*self._num_anchors, axis=1).reshape( (-1,2) )
                if len(self.center_cache)<100:
                    self.center_cache[key] = anchor_centers

            pos_inds = np.where(scores>=threshold)[0]
            bboxes = distance2bbox(anchor_centers, bbox_preds)
            pos_scores = scores[pos_inds]
            pos_bboxes = bboxes[pos_inds]
            scores_list.append(pos_scores)
            bboxes_list.append(pos_bboxes)
            if self.use_kps:
                kpss = distance2kps(anchor_centers, kps_preds)
                #kpss = kps_preds
                kpss = kpss.reshape( (kpss.shape[0], -1, 2) )
                pos_kpss = kpss[pos_inds]
                kpss_list.append(pos_kpss)
        return scores_list, bboxes_list, kpss_list

    def detect(self, img, input_size = None, thresh=None, max_num=0, metric='default'):
        assert input_size is not None or self.input_size is not None
        input_size = self.input_size if input_size is None else input_size
            
        im_ratio = float(img.shape[0]) / img.shape[1]
        model_ratio = float(input_size[1]) / input_size[0]
        if im_ratio>model_ratio:
            new_height = input_size[1]
            new_width = int(new_height / im_ratio)
        else:
            new_width = input_size[0]
            new_height = int(new_width * im_ratio)
        det_scale = float(new_height) / img.shape[0]
        resized_img = cv2.resize(img, (new_width, new_height))
        det_img = np.zeros( (input_size[1], input_size[0], 3), dtype=np.uint8 )
        det_img[:new_height, :new_width, :] = resized_img
        det_thresh = thresh if thresh is not None else self.det_thresh

        scores_list, bboxes_list, kpss_list = self.forward(det_img, det_thresh)

        scores = np.vstack(scores_list)
        scores_ravel = scores.ravel()
        order = scores_ravel.argsort()[::-1]
        bboxes = np.vstack(bboxes_list) / det_scale
        if self.use_kps:
            kpss = np.vstack(kpss_list) / det_scale
        pre_det = np.hstack((bboxes, scores)).astype(np.float32, copy=False)
        pre_det = pre_det[order, :]
        keep = self.nms(pre_det)
        det = pre_det[keep, :]
        if self.use_kps:
            kpss = kpss[order,:,:]
            kpss = kpss[keep,:,:]
        else:
            kpss = None
        if max_num > 0 and det.shape[0] > max_num:
            area = (det[:, 2] - det[:, 0]) * (det[:, 3] -
                                                    det[:, 1])
            img_center = img.shape[0] // 2, img.shape[1] // 2
            offsets = np.vstack([
                (det[:, 0] + det[:, 2]) / 2 - img_center[1],
                (det[:, 1] + det[:, 3]) / 2 - img_center[0]
            ])
            offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
            if metric=='max':
                values = area
            else:
                values = area - offset_dist_squared * 2.0  # some extra weight on the centering
            bindex = np.argsort(
                values)[::-1]  # some extra weight on the centering
            bindex = bindex[0:max_num]
            det = det[bindex, :]
            if kpss is not None:
                kpss = kpss[bindex, :]
        return det, kpss

    def autodetect(self, img, max_num=0, metric='max'):
        bboxes, kpss = self.detect(img, input_size=(640, 640), thresh=0.5)
        bboxes2, kpss2 = self.detect(img, input_size=(128, 128), thresh=0.5)
        bboxes_all = np.concatenate([bboxes, bboxes2], axis=0)
        kpss_all = np.concatenate([kpss, kpss2], axis=0)
        keep = self.nms(bboxes_all)
        det = bboxes_all[keep,:]
        kpss = kpss_all[keep,:]
        if max_num > 0 and det.shape[0] > max_num:
            area = (det[:, 2] - det[:, 0]) * (det[:, 3] -
                                                    det[:, 1])
            img_center = img.shape[0] // 2, img.shape[1] // 2
            offsets = np.vstack([
                (det[:, 0] + det[:, 2]) / 2 - img_center[1],
                (det[:, 1] + det[:, 3]) / 2 - img_center[0]
            ])
            offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
            if metric=='max':
                values = area
            else:
                values = area - offset_dist_squared * 2.0  # some extra weight on the centering
            bindex = np.argsort(
                values)[::-1]  # some extra weight on the centering
            bindex = bindex[0:max_num]
            det = det[bindex, :]
            if kpss is not None:
                kpss = kpss[bindex, :]
        return det, kpss

    def nms(self, dets):
        thresh = self.nms_thresh
        x1 = dets[:, 0]
        y1 = dets[:, 1]
        x2 = dets[:, 2]
        y2 = dets[:, 3]
        scores = dets[:, 4]

        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]

        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)

            inds = np.where(ovr <= thresh)[0]
            order = order[inds + 1]

        return keep

检测结果

为数组item代表每个人脸
confidenceLevel是可信度
其他的是人脸位置

[{'topLeftY': 1663.71, 'topLeftX': 2046.11, 'lowerRightY': 1980.76, 'lowerRightX': 2327.76, 'confidenceLevel': 0.91}, {'topLeftY': 993.82, 'topLeftX': 1126.95, 'lowerRightY': 1341.06, 'lowerRightX': 1409.23, 'confidenceLevel': 0.9}, {'topLeftY': 1249.68, 'topLeftX': 3530.44, 'lowerRightY': 1577.46, 'lowerRightX': 3805.19, 'confidenceLevel': 0.89}, {'topLeftY': 1544.73, 'topLeftX': 2624.07, 'lowerRightY': 1866.22, 'lowerRightX': 2912.39, 'confidenceLevel': 0.88}, {'topLeftY': 737.86, 'topLeftX': 2330.98, 'lowerRightY': 1068.18, 'lowerRightX': 2581.18, 'confidenceLevel': 0.87}]

模型文件

在绑定的资源中