使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(二)——Python代码解读篇

目录

前言

一、padim算法onnx模型输入输出解读

二、padim算法Python代码处理流程分析

2.1 预处理部分

2.2 预测部分

2.3 后处理部分

2.4 可视化部分

三、总结与展望

前言

        上一篇博客中完成了Anomalib中padim算法的模型训练,得到了onnx模型以及推理的效果,想看这部分的同学可以上翻...对于像我一样根本没读论文的同学,获得了onnx模型以后大概率一脸懵,输入是什么?输出是什么?需要经过什么样的预处理和后处理?如何画出和Anomalib项目中一样好看的概率热图呢?C++中如何部署?本篇博客会带大家逐个分析这些问题。本来想和C++部署一起写的,但是实在太长了。想直接看C++代码的同学略过本篇(不过还是建议看一下),下一篇三天内发出来Orz...

一、padim算法onnx模型输入输出解读

        当我们不知道模型里面的结构时,借助Netron工具往往是比较好的办法,Netron网站地址:

Netronhttps://netron.app/        将模型拖入,界面上就会显示网络结构,以下是padim的onnx模型的输入和输出部分的结构:

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(二)——Python代码解读篇_第1张图片使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(二)——Python代码解读篇_第2张图片

        可以看出,其输入尺寸是1*3*256*256,熟悉深度学习的同学应该知道这是一个张量(Tensor),它来源于一张3通道RGB图像,其长宽均为256像素。

        结论1:输入是预处理后的256*256的RGB图像。

        输出也是一个张量,只不过其尺寸为1*1*256*256,见到和原图长宽尺寸相等的数据,我们可以大胆猜测:输出就是我们想要的概率热图,或者某种每个像素位置的得分图,只不过不一定是最终形式。

        结论2:输出是某种得分图,经过后处理后也许可以得到概率热图。

        总的来说,这个模型的输入和输出并不复杂,这对我们进行算法部署是一件好事。

二、padim算法Python代码处理流程分析

        我们虽然急于脱离开Anomalib这个复杂的项目环境,投入下一步的C++部署。但在此之前,我们必须把Python代码运行全过程搞明白,才有可能完成C++的改写。这个过程包括预处理、推理、后处理和可视化四部分。

        由于使用了onnx模型进行推理,根据官网教程,此处应该使用tools/inference/openvino_inference.py进行推理。打开该py文件,很容易在infer函数中找到以下代码段:

    for filename in filenames:
        image = read_image(filename, (256, 256))
        predictions = inferencer.predict(image=image)
        output = visualizer.visualize_image(predictions)

        这几行代码就是我们探究推理过程的根源。首先使用read_image读入图片,然后调用inferencer的predict方法得到推理结果,最后使用visualizer将推理结果可视化。以上就是使用C++部署时要还原的过程。

        read_image比较简单,大家可以自行阅读源码。这里从inferencer开始看起。这里的inferencer由OpenVINOInferencer实例化而来,阅读其predict方法,可以找到如下代码:

        processed_image = self.pre_process(image_arr)              # 预处理
        predictions = self.forward(processed_image)                # 预测   
        output = self.post_process(predictions, metadata=metadata) # 后处理

        return ImageResult(
            image=image_arr,
            pred_score=output["pred_score"],
            pred_label=output["pred_label"],
            anomaly_map=output["anomaly_map"],
            pred_mask=output["pred_mask"],
            pred_boxes=output["pred_boxes"],
            box_labels=output["box_labels"],                       # 返回的各项参数

        可以看到,图像经过了预处理、预测和后处理之后得到了output结果(字典),返回值为各项参数。根据返回值各项参数的名称,我们可以知道该函数返回了预测的得分、标签类别、异常图等等。

2.1 预处理部分

        首先看pre_process部分,在openvino_inferencer.py中定义了该方法:

    def pre_process(self, image: np.ndarray) -> np.ndarray:
        """Pre process the input image by applying transformations.

        Args:
            image (np.ndarray): Input image.

        Returns:
            np.ndarray: pre-processed image.
        """
        transform = A.from_dict(self.metadata["transform"])
        processed_image = transform(image=image)["image"]

        if len(processed_image.shape) == 3:
            processed_image = np.expand_dims(processed_image, axis=0)

        if processed_image.shape[-1] == 3:
            processed_image = processed_image.transpose(0, 3, 1, 2)

        return processed_image

        该方法的核心在于transform,通过查看metadata部分代码,可以认为图片经过了与ImageNet相同的标准化预处理,即RGB的均值为[0.406, 0.456, 0.485],方差为[0.225, 0.224, 0.229]。标准化后需要将其尺寸按照Python代码所示进行修改。

        结论3:预处理步骤包括按照ImageNet标准化处理图像,和处理图像的尺寸。

2.2 预测部分

        其次看forward部分,就在pre_process部分下方:

    def forward(self, image: np.ndarray) -> np.ndarray:
        """Forward-Pass input tensor to the model.

        Args:
            image (np.ndarray): Input tensor.

        Returns:
            np.ndarray: Output predictions.
        """
        return self.network.infer(inputs={self.input_blob: image})

        这部分就是将经过了pre_process的图像送入模型进行预测,很好理解。上篇博客说过此处不纠结神经网络内部的原理,只需要将其当作黑盒使用即可。经过调试查看,发现其输出确实为与原图尺寸相等的得分图,代表了每个像素位置的分数,分数越高越有可能是异常区域。

2.3 后处理部分

        然后是后处理部分,这部分是四部分中最复杂的。

    def post_process(self, predictions: np.ndarray, metadata: dict | DictConfig | None = None) -> dict[str, Any]:
        """Post process the output predictions.

        Args:
            predictions (np.ndarray): Raw output predicted by the model.
            metadata (Dict, optional): Meta data. Post-processing step sometimes requires
                additional meta data such as image shape. This variable comprises such info.
                Defaults to None.

        Returns:
            dict[str, Any]: Post processed prediction results.
        """
        if metadata is None:
            metadata = self.metadata

        predictions = predictions[self.output_blob]

        # Initialize the result variables.
        anomaly_map: np.ndarray | None = None
        pred_label: float | None = None
        pred_mask: float | None = None

        # If predictions returns a single value, this means that the task is
        # classification, and the value is the classification prediction score.
        if len(predictions.shape) == 1:
            task = TaskType.CLASSIFICATION
            pred_score = predictions
        else:
            task = TaskType.SEGMENTATION
            anomaly_map = predictions.squeeze()
            pred_score = anomaly_map.reshape(-1).max()

        # Common practice in anomaly detection is to assign anomalous
        # label to the prediction if the prediction score is greater
        # than the image threshold.
        if "image_threshold" in metadata:
            pred_label = pred_score >= metadata["image_threshold"]

        if task == TaskType.CLASSIFICATION:
            _, pred_score = self._normalize(pred_scores=pred_score, metadata=metadata)
        elif task in (TaskType.SEGMENTATION, TaskType.DETECTION):
            if "pixel_threshold" in metadata:
                pred_mask = (anomaly_map >= metadata["pixel_threshold"]).astype(np.uint8)

            anomaly_map, pred_score = self._normalize(
                pred_scores=pred_score, anomaly_maps=anomaly_map, metadata=metadata
            )
            assert anomaly_map is not None

            if "image_shape" in metadata and anomaly_map.shape != metadata["image_shape"]:
                image_height = metadata["image_shape"][0]
                image_width = metadata["image_shape"][1]
                anomaly_map = cv2.resize(anomaly_map, (image_width, image_height))

                if pred_mask is not None:
                    pred_mask = cv2.resize(pred_mask, (image_width, image_height))
        else:
            raise ValueError(f"Unknown task type: {task}")

        if self.task == TaskType.DETECTION:
            pred_boxes = self._get_boxes(pred_mask)
            box_labels = np.ones(pred_boxes.shape[0])
        else:
            pred_boxes = None
            box_labels = None

        return {
            "anomaly_map": anomaly_map,
            "pred_label": pred_label,
            "pred_score": pred_score,
            "pred_mask": pred_mask,
            "pred_boxes": pred_boxes,
            "box_labels": box_labels,
        }

        上篇博客提到,由于我们使用自制数据集,所以task为classification,所以一切TaskType为SEGMETATION和DETECTION的代码段都不用管。本部分的代码可以精简很多:

    def post_process(self, predictions: np.ndarray, metadata: dict | DictConfig | None = None) -> dict[str, Any]:
        """Post process the output predictions.

        Args:
            predictions (np.ndarray): Raw output predicted by the model.
            metadata (Dict, optional): Meta data. Post-processing step sometimes requires
                additional meta data such as image shape. This variable comprises such info.
                Defaults to None.

        Returns:
            dict[str, Any]: Post processed prediction results.
        """
        if metadata is None:
            metadata = self.metadata

        predictions = predictions[self.output_blob]

        # Initialize the result variables.
        anomaly_map: np.ndarray | None = None
        pred_label: float | None = None
        pred_mask: float | None = None

        # If predictions returns a single value, this means that the task is
        # classification, and the value is the classification prediction score.
        if len(predictions.shape) == 1:
            task = TaskType.CLASSIFICATION
            pred_score = predictions

        # Common practice in anomaly detection is to assign anomalous
        # label to the prediction if the prediction score is greater
        # than the image threshold.
        if "image_threshold" in metadata:
            pred_label = pred_score >= metadata["image_threshold"]

        if task == TaskType.CLASSIFICATION:
            _, pred_score = self._normalize(pred_scores=pred_score, metadata=metadata)
        
        pred_boxes = None
        box_labels = None

        return {
            "anomaly_map": anomaly_map,
            "pred_label": pred_label,
            "pred_score": pred_score,
            "pred_mask": pred_mask,
            "pred_boxes": pred_boxes,
            "box_labels": box_labels,
        }

        阅读源码,发现本部分有意义的输出只有pred_score,真正的处理步骤只有一行:

_, pred_score = self._normalize(pred_scores=pred_score, metadata=metadata)

        进入_normalize部分,可以看到输入的pred_scores是一个张量。事实上,pred_scores即为和原图尺寸相等的概率得分图。同样,输入的pred_scores也只处理了一步,即:

            pred_scores = normalize_min_max(
                pred_scores,
                metadata["image_threshold"],
                metadata["min"],
                metadata["max"],
            )

        再进入normalize_min_max部分,可以看到该函数对pred_scores进行了如下处理:

normalized = ((targets - threshold) / (max_val - min_val)) + 0.5

        这里的max_val和min_val来源于哪里呢?打开训练结果文件夹results/padim/tube/run/weights/onnx/metadata.json,可以文件末尾看到如下信息(tube是我自己的数据集名字):

    "image_threshold": 13.702226638793945,
    "pixel_threshold": 13.702226638793945,
    "min": 5.296699047088623,
    "max": 22.767864227294922

        min即为min_val,max即为max_val,其含义分别为pred_score得分图中的最小值和最大值,image_threshold是计算得到的阈值,得分图中大于该阈值的像素位置我们认为它属于异常区域(缺陷),小于该阈值的区域认为是正常区域。经过这一步标准化以后输出pred_scores。

        结论4:后处理部分输入为预测部分得到的得分图,输出为标准化后的pred_scores。

2.4 可视化部分

        至此,数据处理部分结束,接下来就是如何将数据以概率热图的方式可视化的问题了。回到openvino_inference.py,可以看到visualizer调用了visualize_image方法对数据结果predictions进行了处理,并使用show方法进行可视化。

    for filename in filenames:
        image = read_image(filename, (256, 256))
        predictions = inferencer.predict(image=image)
        output = visualizer.visualize_image(predictions)

        if args.output is None and args.show is False:
            warnings.warn(
                "Neither output path is provided nor show flag is set. Inferencer will run but return nothing."
            )

        if args.output:
            file_path = generate_output_image_filename(input_path=filename, output_path=args.output)
            visualizer.save(file_path=file_path, image=output)

        # Show the image in case the flag is set by the user.
        if args.show:
            visualizer.show(title="Output Image", image=output)

         进入visualize_image方法,我们之前在config.yaml文件中将显示模式设为full,visualize_image方法内使用的是_visualize_full方法。

        if self.mode == "full":
            return self._visualize_full(image_result)

        层层递进,进入_visualize_full方法,同上,只需要关注task为CLASSIFICATION的代码段。在_visualize_full方法中,可以看到如下代码:

        elif self.task == TaskType.CLASSIFICATION:
            visualization.add_image(image_result.image, title="Image")
            if hasattr(image_result, "heat_map"):
                visualization.add_image(image_result.heat_map, "Predicted Heat Map")
            if image_result.pred_label:
                image_classified = add_anomalous_label(image_result.image, image_result.pred_score)
            else:
                image_classified = add_normal_label(image_result.image, 1 - image_result.pred_score)
            visualization.add_image(image=image_classified, title="Prediction")

        visualization.add_image的作用实际上就是往results/padim/tube/run/images的结果中添加图片,添加的三张图片分别为“Image”“Predicted Heat Map”“Prediction”,恰好对应输出的1*3结果图:

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(二)——Python代码解读篇_第3张图片

         这里最关注的还是Predicted Heat Map是怎么画出来的。进入image_result.heat_map,发现它是调用了superimpose_anomaly_map函数生成的:

self.heat_map = superimpose_anomaly_map(self.anomaly_map, self.image, normalize=False)

         再次进入superimpose_anomaly_map函数,其代码段如下:

def superimpose_anomaly_map(
    anomaly_map: np.ndarray, image: np.ndarray, alpha: float = 0.4, gamma: int = 0, normalize: bool = False
) -> np.ndarray:
    """Superimpose anomaly map on top of in the input image.

    Args:
        anomaly_map (np.ndarray): Anomaly map
        image (np.ndarray): Input image
        alpha (float, optional): Weight to overlay anomaly map
            on the input image. Defaults to 0.4.
        gamma (int, optional): Value to add to the blended image
            to smooth the processing. Defaults to 0. Overall,
            the formula to compute the blended image is
            I' = (alpha*I1 + (1-alpha)*I2) + gamma
        normalize: whether or not the anomaly maps should
            be normalized to image min-max


    Returns:
        np.ndarray: Image with anomaly map superimposed on top of it.
    """

    anomaly_map = anomaly_map_to_color_map(anomaly_map.squeeze(), normalize=normalize)
    superimposed_map = cv2.addWeighted(anomaly_map, alpha, image, (1 - alpha), gamma)
    return superimposed_map

        这里的英文注释写的很好了,其实概率热图就是将未处理的原图像和处理后的anomaly_map进行了一定的权重叠加,绘制出的图像。叠加之前,需要对输入的anomaly_map进行anomaly_map_to_color_map函数的处理,而anomaly_map_to_color_map就是将anomaly_map转化为uint8格式的灰度图(像素值0-255),然后根据灰度值绘制伪彩色图:

anomaly_map = cv2.applyColorMap(anomaly_map, cv2.COLORMAP_JET)

        叠加处理后,一幅清晰明艳的缺陷概率热图就产生了。

        结论5:可视化过程中,概率热图为原图和伪彩色anomaly_map的叠加图像。

        到这里,我们已经在Python代码的解读中拉通了整个流程:从输入图像到预处理、预测、后处理和可视化,了解了概率热图的绘制方法。

三、总结与展望

        若想在C++中部署模型,本篇博客这样耗时又繁琐的代码阅读过程是少不了的。只有先明白原工程中Python代码的逻辑,剥离完整的流程,才可能在C++中进行复现。下一篇博客将以C++代码为主,讲解如何使用OnnxRuntime引擎完成模型的部署。感谢阅读和关注~

你可能感兴趣的:(python,视觉检测,人工智能,深度学习)