昇腾AI行业案例（二）：基于 AI 图像处理的安全帽检测

# DO NOT CHANGE THIS CELL
# activate this cell by selecting it with the mouse or arrow keys then use the keyboard shortcut [Shift+Enter] to execute
print('This is just a simple print statement')
This is just a simple print statement

import cv2
import glob
import numpy as np
from pathlib import Path

class LoadImages:
    def __init__(self, path, img_size=640, stride=32, auto=True):
        """
        Initializes YOLOv5 loader for images, supporting directories.
        Args:
            path (str): Directory path containing images.
            img_size (int): Image size for resizing.
            stride (int): Stride for resizing.
            auto (bool): Automatic image resizing.
        """
        self.img_size = img_size
        self.stride = stride
        self.auto = auto
        self.files = sorted(glob.glob(str(Path(path).resolve()) + '/*.*'))  # directory
        self.nf = len(self.files)  # number of files
        assert self.nf > 0, f"No images found in {path}. Supported formats are: {IMG_FORMATS}"
    def __iter__(self):
        """Initializes iterator by resetting count and returns the iterator object itself."""
        self.count = 0
        return self
    def __next__(self):
        """Advances to the next file in the dataset, raising StopIteration if at the end."""
        if self.count == self.nf:
            raise StopIteration
        path = self.files[self.count]
        self.count += 1
        im0 = cv2.imread(path)  # BGR
        assert im0 is not None, f"Image Not Found {path}"
        # Padded resize
        im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0]  # padded resize
        im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        im = np.ascontiguousarray(im)  # contiguous
        return path, im, im0

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    """Resizes and pads image to new_shape with stride-multiple constraints, returns resized image, ratio, padding."""
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)
    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
    dw /= 2  # divide padding into 2 sides
    dh /= 2
    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)

source = "./test_data/"
dataset = LoadImages(source)
for path, im, im0 in dataset:
    print(im.shape)
(3, 448, 640)
(3, 640, 416)

# git clone https://github.com/ultralytics/yolov5
# export PYTHONPATH="./yolov5/:$PYTHONPATH"

from models.common import DetectMultiBackend  # 导入yolov5_s模型
import torch
import torch_npu  # 使用昇腾硬件推理
from torch_npu.contrib import transfer_to_npu  # 使用昇腾硬件推理
import warnings
warnings.filterwarnings('ignore')
model = DetectMultiBackend("./helmet_head_person_s.pt")
img_loader = LoadImages('./test_data')
for path, im, im0s in img_loader:
    im = torch.from_numpy(im)
    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
    im /= 255  # 0 - 255 to 0.0 - 1.0
    if len(im.shape) == 3:
        im = im[None]  # expand for batch dim
    pred = model(im)[0]
    print(pred.shape)
Fusing layers... 
YOLOv5s summary: 232 layers, 7251912 parameters, 0 gradients
torch.Size([1, 17640, 8])
torch.Size([1, 16380, 8])

# (x, y, w, h)的坐标格式转成(x1, y1, x2, y2)
def xywh2xyxy(x):
    """Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right."""
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
    return y
# 计算方框之间的 IOU（Intersection over Union）
def box_iou(box1, box2, eps=1e-7):
    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
    """
    Return intersection-over-union (Jaccard index) of boxes.
    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
    Arguments:
        box1 (Tensor[N, 4])
        box2 (Tensor[M, 4])
    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise
            IoU values for every element in boxes1 and boxes2
    """
    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
    (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
    inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
    # IoU = inter / (area1 + area2 - inter)
    return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)

# 非极大抑制去除重复框
def non_max_suppression(
    prediction,
    conf_thres=0.25,
    iou_thres=0.45,
    classes=None,
    agnostic=False,
    multi_label=True,
    labels=(),
    max_det=300,
    nm=0,  # 掩码的数量
):
    """
    在推理结果上执行非最大抑制（NMS）以拒绝重叠的检测结果。
    返回：
         检测结果的列表，每个图像对应一个(n,6)张量 [x1, y1, x2, y2, conf, cls]
    """
    # 入参检查
    assert 0 <= conf_thres <= 1, f"置信度阈值 {conf_thres} 无效，有效值应在0.0和1.0之间"
    assert 0 <= iou_thres <= 1, f"IoU {iou_thres} 无效，有效值应在0.0和1.0之间"
    if isinstance(prediction, (list, tuple)):  # YOLOv5模型在验证模型中，输出 = (推理输出, 损失输出)
        prediction = prediction[0]  # 仅选择推理输出
    bs = prediction.shape[0]  # 批量大小
    nc = prediction.shape[2] - nm - 5  # 类别数量
    xc = prediction[..., 4] > conf_thres  # 候选框
    # 设置超参
    max_wh = 7680  # （像素）最大框宽和高
    max_nms = 30000  # 传入torchvision.ops.nms()的最大框数
    time_limit = 0.5 + 0.05 * bs  # 超时时间
    redundant = True  # 需要冗余检测
    multi_label &= nc > 1  # 每个框多个标签（增加0.5ms/图像）
    merge = False  # 不使用合并NMS
    t = time.time()
    mi = 5 + nc  # 掩码起始索引
    output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs  # 初始化输出检测框
    for xi, x in enumerate(prediction):  # 遍历推理检测结果
        x = x[xc[xi]]  # 获取置信度
        # 如果自动标记，则添加先验标签
        if labels and len(labels[xi]):
            lb = labels[xi]
            v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
            v[:, :4] = lb[:, 1:5]  # 框
            v[:, 4] = 1.0  # 置信度
            v[range(len(lb)), lb[:, 0].long() + 5] = 1.0  # 类别
            x = torch.cat((x, v), 0)
        # 如果没有剩余的处理下一张图像
        if not x.shape[0]:
            continue
        # 计算置信度
        x[:, 5:] *= x[:, 4:5]  # 置信度 = 目标置信度 * 类别置信度
        # 框/掩码
        box = xywh2xyxy(x[:, :4])  # (中心x, 中心y, 宽度, 高度) 转换为 (x1, y1, x2, y2)
        mask = x[:, mi:]  # 如果没有掩码则为零列
        # 检测矩阵 nx6 (xyxy, conf, cls)
        if multi_label:
            i, j = (x[:, 5:mi] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)
        else:  # 仅最佳类别
            conf, j = x[:, 5:mi].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
        # 按类别过滤
        if classes is not None:
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
        # 检查形状
        n = x.shape[0]  # 框的数量
        if not n:  # 没有框
            continue
        x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # 按置信度排序并移除多余的框
        # 批量NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # 类别
        boxes, scores = x[:, :4] + c, x[:, 4]  # 框（按类别偏移）, 得分
        # 调用torchvision封装好的NMS接口
        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
        i = i[:max_det]  # 限制检测数量
        if merge and (1 < n < 3e3):  # 合并NMS（使用加权平均合并框）
            # 更新框为 boxes(i,4) = weights(i,n) * boxes(n,4)
            iou = box_iou(boxes[i], boxes) > iou_thres  # IoU矩阵
            weights = iou * scores[None]  # 框权重
            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # 合并框
            if redundant:
                i = i[iou.sum(1) > 1]  # 需要冗余
        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            LOGGER.warning(f"警告 ⚠️ NMS 时间限制 {time_limit:.3f}s 超过")
            break  # 时间限制超过
    return output

def clip_boxes(boxes, shape):
    """Clips bounding box coordinates (xyxy) to fit within the specified image shape (height, width)."""
    if isinstance(boxes, torch.Tensor):  # faster individually
        boxes[..., 0].clamp_(0, shape[1])  # x1
        boxes[..., 1].clamp_(0, shape[0])  # y1
        boxes[..., 2].clamp_(0, shape[1])  # x2
        boxes[..., 3].clamp_(0, shape[0])  # y2
    else:  # np.array (faster grouped)
        boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1])  # x1, x2
        boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0])  # y1, y2
def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
    """Rescales (xyxy) bounding boxes from img1_shape to img0_shape, optionally using provided `ratio_pad`."""
    if ratio_pad is None:  # calculate from img0_shape
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]
    boxes[..., [0, 2]] -= pad[0]  # x padding
    boxes[..., [1, 3]] -= pad[1]  # y padding
    boxes[..., :4] /= gain
    clip_boxes(boxes, img0_shape)
    return boxes

# 找到没有佩戴安全帽的人
def find_person_without_helmet(detections):
    person_without_helmet = []
    # 将检测结果按类别分开存储
    person_detections = []
    helmet_detections = []
    for det in detections:
        x1, y1, x2, y2, confidence, label = det
        if label == 0:  # 'person'
            person_detections.append((x1, y1, x2, y2))
        elif label == 2:  # 'helmet'
            helmet_detections.append((x1, y1, x2, y2))
    # 检查每个'person'检测框是否包含'helmet'
    for person_box in person_detections:
        person_x1, person_y1, person_x2, person_y2 = person_box
        person_area = (person_x2 - person_x1) * (person_y2 - person_y1)
        # 检查是否有'helmet'的90%区域落在'person'检测框内
        contains_helmet = False
        for helmet_box in helmet_detections:
            helmet_x1, helmet_y1, helmet_x2, helmet_y2 = helmet_box
            helmet_area = (helmet_x2 - helmet_x1) * (helmet_y2 - helmet_y1)
            # 计算'helmet'和'person'的交集面积
            intersection_x1 = max(person_x1, helmet_x1)
            intersection_y1 = max(person_y1, helmet_y1)
            intersection_x2 = min(person_x2, helmet_x2)
            intersection_y2 = min(person_y2, helmet_y2)
            intersection_area = max(0, intersection_x2 - intersection_x1) * max(0, intersection_y2 - intersection_y1)
            # 检查'helmet'的90%是否在'person'内
            if intersection_area / helmet_area >= 0.9:
                contains_helmet = True
                break
        # 如果当前'person'不包含'helmet'
        if not contains_helmet:
            person_without_helmet.append((person_x1, person_y1, person_x2, person_y2))
    return person_without_helmet

from ultralytics.utils.plotting import Annotator, colors
import time
import torchvision
conf_thres=0.25  # confidence threshold
iou_thres=0.45  # NMS IOU threshold
max_det=1000  # maximum detections per image
hide_labels = False
hide_conf = False
line_thickness = 3
model = DetectMultiBackend("./helmet_head_person_s.pt")
stride, names, pt = model.stride, model.names, model.pt
imgsz=(640, 640)
model.warmup(imgsz=(1 if pt or model.triton else 1, 3, *imgsz))  # warmup
img_loader = LoadImages('./test_data')
save_dir = "./detect_imgs/"
# Run inference
for path, im, im0s in img_loader:
    im = torch.from_numpy(im)
    im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
    im /= 255  # 0 - 255 to 0.0 - 1.0
    if len(im.shape) == 3:
        im = im[None]  # expand for batch dim
    pred = model(im)[0]
    # Apply NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False, max_det=max_det)
    # Process detections
    for i, det in enumerate(pred):  # per image
        annotator = Annotator(im0s, line_width=line_thickness, example=str(names))
        if len(det):
            number_of_person = np.sum(np.array(det[:, 5] == 0.0))
            number_of_helmet = np.sum(np.array(det[:, 5] == 2.0))
            if number_of_helmet < number_of_person:
                print("Notice! Someone is not wearing a helmet！")
            no_helmet_person = find_person_without_helmet(det)
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0s.shape).round()
            # Print results
            for *xyxy, conf, cls in reversed(det):
                c = int(cls)  # integer class
                if tuple(xyxy) in no_helmet_person:
                    label = "No helmet, Dangerous!"
                    annotator.box_label(xyxy, label, color=(0, 0, 255))
                else:
                    label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                    annotator.box_label(xyxy, label, color=colors(c, True))
    # Save results (image with detections)
    print("save_path: ", save_dir + path.split("/")[-1])
    cv2.imwrite(str(save_dir + path.split("/")[-1]), im0s)
Fusing layers... 
YOLOv5s summary: 232 layers, 7251912 parameters, 0 gradients
save_path:  ./detect_imgs/000012.jpg
Notice! Someone is not wearing a helmet！
save_path:  ./detect_imgs/000178.jpg