昇腾 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_images/"dataset = LoadImages(source)count = 0for path, im, im0 in dataset:    print(im.shape)    if count > 5:        break    count += 1

(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)(3, 480, 640)

# 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')  # 屏蔽warningsmodel = DetectMultiBackend("./best.pt")  # 实例化模型img_loader = LoadImages('./test_images')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, 1)    if len(im.shape) == 3:        im = im[None]  # expand for batch dim    pred = model(im)[0]  # 调用yolo进行推理    print(pred.shape)    break

Fusing layers... YOLOv5m summary: 290 layers, 20889303 parameters, 0 gradients, 48.0 GFLOPs

torch.Size([1, 18900, 15])

# 把检测框的[x, y, w, h]坐标转成[x1, y1, x2, y2]坐标def xywh2xyxy(x):    """[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
# 计算检测框之前的重合度（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

conf_thres=0.25  # 置信度阈值iou_thres=0.45  # NMS IOU 阈值max_det=1000  # 最大的检测框数量imgsz=(640, 640)  # 图片目标形状save_dir = "./detect_imgs/"  # 检测结果保存目录

import timeimport torchvisionfrom ultralytics.utils.plotting import Annotator, colors
model = DetectMultiBackend("./best.pt")names = model.namesprint("classification names: ", names)img_loader = LoadImages('./test_images')# Run inferencefor path, im, im0s in img_loader:    # 把输入处理成AI模型的输入格式    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]
    # NMS 后处理    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False, max_det=max_det)    print("%s defects were found in the aluminum plate." % pred[0].shape[0])        # 把检测结果画在原图上    for i, det in enumerate(pred):  # per image        annotator = Annotator(im0s, line_width=8, example=str(names))        if len(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                label = f'{names[c]} {conf:.2f}'                annotator.box_label(xyxy, label, color=colors(c, True))
    # 保存检测结果    print("save_path: ", save_dir + path.split("/")[-1])    cv2.imwrite(str(save_dir + path.split("/")[-1]), im0s)

Fusing layers... YOLOv5m summary: 290 layers, 20889303 parameters, 0 gradients, 48.0 GFLOPs

classification names:  {0: 'non_conduct', 1: 'abrasion_mark', 2: 'corner_leak', 3: 'orange_peel', 4: 'leak', 5: 'jet_flow', 6: 'paint_bubble', 7: 'pit', 8: 'motley', 9: 'dirty_spot'}1 defects were found in the aluminum plate.save_path:  ./detect_imgs/000015.jpg1 defects were found in the aluminum plate.save_path:  ./detect_imgs/000427.jpg1 defects were found in the aluminum plate.save_path:  ./detect_imgs/000467.jpg1 defects were found in the aluminum plate.save_path:  ./detect_imgs/000490.jpg1 defects were found in the aluminum plate.save_path:  ./detect_imgs/000978.jpg2 defects were found in the aluminum plate.save_path:  ./detect_imgs/001178.jpg1 defects were found in the aluminum plate.save_path:  ./detect_imgs/002026.jpg