昇腾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_images/"
dataset = LoadImages(source)
count = 0
for 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 torch
import torch_npu  # 使用昇腾硬件推理
from torch_npu.contrib import transfer_to_npu  # 使用昇腾硬件推理
import warnings
warnings.filterwarnings('ignore')  # 屏蔽warnings
model = 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 time
import torchvision
from ultralytics.utils.plotting import Annotator, colors
model = DetectMultiBackend("./best.pt")
names = model.names
print("classification names: ", names)
img_loader = LoadImages('./test_images')
# Run inference
for 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.jpg
1 defects were found in the aluminum plate.
save_path:  ./detect_imgs/000427.jpg
1 defects were found in the aluminum plate.
save_path:  ./detect_imgs/000467.jpg
1 defects were found in the aluminum plate.
save_path:  ./detect_imgs/000490.jpg
1 defects were found in the aluminum plate.
save_path:  ./detect_imgs/000978.jpg
2 defects were found in the aluminum plate.
save_path:  ./detect_imgs/001178.jpg
1 defects were found in the aluminum plate.
save_path:  ./detect_imgs/002026.jpg