昇腾 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 executeprint('This is just a simple print statement')

This is just a simple print statement

import cv2import globimport numpy as npfrom 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 torchimport torch_npu  # 使用昇腾硬件推理from torch_npu.contrib import transfer_to_npu  # 使用昇腾硬件推理import warningswarnings.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, colorsimport timeimport torchvision
conf_thres=0.25  # confidence thresholdiou_thres=0.45  # NMS IOU thresholdmax_det=1000  # maximum detections per imagehide_labels = Falsehide_conf = Falseline_thickness = 3model = DetectMultiBackend("./helmet_head_person_s.pt")stride, names, pt = model.stride, model.names, model.ptimgsz=(640, 640)model.warmup(imgsz=(1 if pt or model.triton else 1, 3, *imgsz))  # warmupimg_loader = LoadImages('./test_data')save_dir = "./detect_imgs/"# Run inferencefor 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.jpgNotice! Someone is not wearing a helmet！save_path:  ./detect_imgs/000178.jpg