How to Use Custom Model for Automatic Annotation?

[haitian-zhang]

Search before asking

• I have searched the X-AnyLabeling model_zoo and found no similar model requests.

Description

你好，我的数据特征和公开特征不太相似，是否能在x-labeling中加入自己训练的模型权重用于半自动标注

Use case

No response

Additional

No response

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[CVHub520]

Hello, thank you for raising this issue. Yes, you can use your own dataset-trained model weights for semi-automatic annotation. Please refer to the custom model section in our documentation for instructions on how to do this: https://github.com/CVHub520/X-AnyLabeling/blob/main/docs/zh_cn/custom_model.md

[haitian-zhang]

我保存的模型文件不是onnx格式的，是否先进行转换，这一步应该如何实现呢

[CVHub520]

Thank you for your question. Suer, you need to convert your model to ONNX format. The custom model section in our documentation already includes instructions on how to convert a PyTorch model to ONNX format. Please refer the above link. You can find detailed conversion steps in this document.

[haitian-zhang]

好嘞，谢谢大佬

[haitian-zhang]

你好，我上传的文件是unet系列的自定义模型，我在我原本的代码里测试了onnx文件可以成功预测出分类结果，但是按照教程中复制了参考示例的unet.py部分放入了软件当成推理文件却出现了报错，是否需要结合我自己的推理文件来修改配置的推理文件，我上传了两个py文件，其中convert_pth_to_onnx.py为我转换onnx和测试推理效果的代码，unetplus,py为我在x-labeling软件中配置的推理文件
convert_pth_to_onnx.py：import os
import cv2
import onnxruntime
import torch
import glob
import numpy as np
from albumentations import Compose
from albumentations.augmentations import transforms
from torch.utils.data import DataLoader
import segmentation_models_pytorch as smp

import L1_archs_cut

from dataset import test_Dataset

from dataProcess import get_dataloader
DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
from torch.optim.swa_utils import AveragedModel
def pth_2onnx():
"""
pytorch 模型转换为onnx模型
:return:
"""
channels = 3
torch_model = torch.load('../user_data/model_data/unetplusplus_dpn132_upsample_SoftCE_dice.pth')
model = smp.UnetPlusPlus(
encoder_name="timm-regnety_320", # timm-resnest101e
# encoder_name="resnet50",
encoder_weights="imagenet",
in_channels=channels,
classes=2,
)
# model = L1_archs_cut.NestedUNet(num_classes=1, input_channels=3, deep_supervision=True)
model.load_state_dict(torch_model)
batch_size = 32 # 批处理大小
input_shape = (3, 256, 256) # 输入数据

# set the model to inference mode
model.eval()
print(model)
x = torch.randn(batch_size, *input_shape)  # 生成张量
export_onnx_file = "../user_data/model_data/Unetplus.onnx"  # 目的ONNX文件名
torch.onnx.export(model,
                  x,
                  export_onnx_file,
                  # 注意这个地方版本选择为11
                  opset_version=11,
                  do_constant_folding=True,  # 是否执行常量折叠优化
                  input_names=["input"],  # 输入名
                  output_names=["output"],  # 输出名
                  dynamic_axes={"input": {0: "batch_size"},  # 批处理变量
                                "output": {0: "batch_size"}})

def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def test_1(model_path, output_dir, test_loader, addNDVI):
in_channels = 3
if(addNDVI):
in_channels += 1
model = smp.UnetPlusPlus(
encoder_name="timm-regnety_320", #resnet101
encoder_weights="imagenet",
in_channels=in_channels,
classes=2,
)
# model = smp.DeepLabV3Plus(
# encoder_name="resnet101",
# encoder_weights="imagenet",
# in_channels=4,
# classes=8,
# )
# 如果模型是SWA
if("swa" in model_path):
model = AveragedModel(model)
model.to(DEVICE);
# load the mistake package
session = onnxruntime.InferenceSession(model_path,providers=['CUDAExecutionProvider' if DEVICE=='cuda' else 'CPUExecutionProvider'])
input_name = session.get_inputs()[0].name
# model.load_state_dict(torch.load(model_path))
# model.eval()

# for image, image_stretch, image_path, ndvi in test_loader:
for image, image_stretch, image_path in test_loader:
    with torch.no_grad():
        # image = image.cuda()
        # image_stretch = image_stretch.cuda()
        image = image.cpu().numpy().astype(np.float32)
        image_stretch = image_stretch.cpu().numpy().astype(np.float32)
        output1 = session.run(None,{input_name: image})[0]
        output2 = session.run(None, {input_name:image_stretch})[0]
        # output1 = model(image).cpu().data.numpy()
        # output2 = model(image_stretch).cpu().data.numpy()
    output = (output1 + output2) / 2.0
    for i in range(output.shape[0]):
        pred = output[i]
        pred = np.argmax(pred, axis = 0) + 1
        pred = np.uint8(pred)
        # save_path = os.path.join(output_dir, image_path[i].split('\\')[-1].replace('.tif', '.png'))
        # save_path = os.path.join(output_dir, image_path[i].split('\\')[-1])
        save_path = os.path.join(output_dir, os.path.basename(image_path[i]))
        print(image_path[i])
        print(save_path)
        cv2.imwrite(save_path, pred)

if name == 'main':
# pth_2onnx()
# pth_2onnx()
# image_test()
batch_size = 32
addNDVI = False
model_path = "../user_data/model_data/Unetplus.onnx"
output_dir = '../prediction_result/pre'
test_image_paths = glob.glob('../tcdata/test/image/*.png')
test_loader = get_dataloader(test_image_paths, None, "test", addNDVI, batch_size, False, 8)
# test(model_path_effi7, model_path_resnest, output_dir, test_loader, addNDVI)
test_1(model_path, output_dir, test_loader, addNDVI)

[haitian-zhang]

这是我在x-labeling配置的推理文件unetplus.py:import logging
import os

import cv2
import numpy as np
from PyQt5 import QtCore
from PyQt5.QtCore import QCoreApplication

from anylabeling.app_info import preferred_device
from anylabeling.views.labeling.shape import Shape
from anylabeling.views.labeling.utils.opencv import qt_img_to_rgb_cv_img
from .model import Model
from .types import AutoLabelingResult
from .engines.build_onnx_engine import OnnxBaseModel

class UNetPlus(Model):
"""Semantic segmentation model using UNet"""

class Meta:
    required_config_names = [
        "type",
        "name",
        "display_name",
        "model_path",
        "classes",
    ]
    widgets = ["button_run"]
    output_modes = {
        "polygon": QCoreApplication.translate("Model", "Polygon"),
    }
    default_output_mode = "polygon"

def __init__(self, model_config, on_message) -> None:
    # Run the parent class's init method
    super().__init__(model_config, on_message)
    model_name = self.config["type"]
    model_abs_path = self.get_model_abs_path(self.config, "model_path")
    if not model_abs_path or not os.path.isfile(model_abs_path):
        raise FileNotFoundError(
            QCoreApplication.translate(
                "Model",
                f"Could not download or initialize {model_name} model.",
            )
        )
    self.net = OnnxBaseModel(model_abs_path, __preferred_device__)
    self.classes = self.config["classes"]
    self.input_shape = self.net.get_input_shape()[-2:]

def preprocess(self, input_image):
    """
    Pre-processes the input image before feeding it to the network.

    Args:
        input_image (numpy.ndarray): The input image to be processed.

    Returns:
        numpy.ndarray: The pre-processed output.
    """
    input_h, input_w = self.input_shape
    image = cv2.resize(input_image, (input_w, input_h))
    image = np.transpose(image, (2, 0, 1))
    image = image.astype(np.float32) / 255.0
    image = (image - 0.5) / 0.5
    image = np.expand_dims(image, axis=0)
    return image

def postprocess(self, image, outputs):
    """
    Post-processes the network's output.

    Args:
        image (numpy.ndarray): The input image.
        outputs (numpy.ndarray): The output from the network.

    Returns:
        contours (list): List of contours for each detected object class.
                        Each contour is represented as a dictionary containing
                        the class label and a list of contour points.
    """
    n, c, h, w = outputs.shape
    image_height, image_width = image.shape[:2]
    # Obtain the category index of each pixel
    # target shape: (1, h, w)
    outputs = np.argmax(outputs, axis=1)
    results = []
    for i in range(c):
        # Skip the background label
        if self.classes[i] == '_background_':
            continue
        # Get the category index of each pixel for the first batch by adding [0].
        mask = outputs[0] == i
        # Rescaled to original shape
        mask_resized = cv2.resize(mask.astype(np.uint8), (image_width, image_height))
        # Get the contours
        contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # Append the contours along with their respective class labels
        results.append((self.classes[i], [np.squeeze(contour).tolist() for contour in contours]))
    return results

def predict_shapes(self, image, image_path=None):
    """
    Predict shapes from image
    """

    if image is None:
        return []

    try:
        image = qt_img_to_rgb_cv_img(image, image_path)
    except Exception as e:  # noqa
        logging.warning("Could not inference model")
        logging.warning(e)
        return []

    blob = self.preprocess(image)
    outputs = self.net.get_ort_inference(blob)
    results = self.postprocess(image, outputs)
    shapes = []
    for item in results:
        label, contours = item
        for points in contours:
            # Make sure to close
            points += points[0]
            shape = Shape(flags={})
            for point in points:
                shape.add_point(QtCore.QPointF(point[0], point[1]))
            shape.shape_type = "polygon"
            shape.closed = True
            shape.fill_color = "#000000"
            shape.line_color = "#000000"
            shape.line_width = 1
            shape.label = label
            shape.selected = False
            shapes.append(shape)

    result = AutoLabelingResult(shapes, replace=True)
    return result

def unload(self):
    del self.net

[CVHub520]

这是我在x-labeling配置的推理文件unetplus.py:import logging import os

import cv2 import numpy as np from PyQt5 import QtCore from PyQt5.QtCore import QCoreApplication

from anylabeling.app_info import preferred_device from anylabeling.views.labeling.shape import Shape from anylabeling.views.labeling.utils.opencv import qt_img_to_rgb_cv_img from .model import Model from .types import AutoLabelingResult from .engines.build_onnx_engine import OnnxBaseModel

class UNetPlus(Model): """Semantic segmentation model using UNet"""

class Meta:
    required_config_names = [
        "type",
        "name",
        "display_name",
        "model_path",
        "classes",
    ]
    widgets = ["button_run"]
    output_modes = {
        "polygon": QCoreApplication.translate("Model", "Polygon"),
    }
    default_output_mode = "polygon"

def __init__(self, model_config, on_message) -> None:
    # Run the parent class's init method
    super().__init__(model_config, on_message)
    model_name = self.config["type"]
    model_abs_path = self.get_model_abs_path(self.config, "model_path")
    if not model_abs_path or not os.path.isfile(model_abs_path):
        raise FileNotFoundError(
            QCoreApplication.translate(
                "Model",
                f"Could not download or initialize {model_name} model.",
            )
        )
    self.net = OnnxBaseModel(model_abs_path, __preferred_device__)
    self.classes = self.config["classes"]
    self.input_shape = self.net.get_input_shape()[-2:]

def preprocess(self, input_image):
    """
    Pre-processes the input image before feeding it to the network.

    Args:
        input_image (numpy.ndarray): The input image to be processed.

    Returns:
        numpy.ndarray: The pre-processed output.
    """
    input_h, input_w = self.input_shape
    image = cv2.resize(input_image, (input_w, input_h))
    image = np.transpose(image, (2, 0, 1))
    image = image.astype(np.float32) / 255.0
    image = (image - 0.5) / 0.5
    image = np.expand_dims(image, axis=0)
    return image

def postprocess(self, image, outputs):
    """
    Post-processes the network's output.

    Args:
        image (numpy.ndarray): The input image.
        outputs (numpy.ndarray): The output from the network.

    Returns:
        contours (list): List of contours for each detected object class.
                        Each contour is represented as a dictionary containing
                        the class label and a list of contour points.
    """
    n, c, h, w = outputs.shape
    image_height, image_width = image.shape[:2]
    # Obtain the category index of each pixel
    # target shape: (1, h, w)
    outputs = np.argmax(outputs, axis=1)
    results = []
    for i in range(c):
        # Skip the background label
        if self.classes[i] == '_background_':
            continue
        # Get the category index of each pixel for the first batch by adding [0].
        mask = outputs[0] == i
        # Rescaled to original shape
        mask_resized = cv2.resize(mask.astype(np.uint8), (image_width, image_height))
        # Get the contours
        contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # Append the contours along with their respective class labels
        results.append((self.classes[i], [np.squeeze(contour).tolist() for contour in contours]))
    return results

def predict_shapes(self, image, image_path=None):
    """
    Predict shapes from image
    """

    if image is None:
        return []

    try:
        image = qt_img_to_rgb_cv_img(image, image_path)
    except Exception as e:  # noqa
        logging.warning("Could not inference model")
        logging.warning(e)
        return []

    blob = self.preprocess(image)
    outputs = self.net.get_ort_inference(blob)
    results = self.postprocess(image, outputs)
    shapes = []
    for item in results:
        label, contours = item
        for points in contours:
            # Make sure to close
            points += points[0]
            shape = Shape(flags={})
            for point in points:
                shape.add_point(QtCore.QPointF(point[0], point[1]))
            shape.shape_type = "polygon"
            shape.closed = True
            shape.fill_color = "#000000"
            shape.line_color = "#000000"
            shape.line_width = 1
            shape.label = label
            shape.selected = False
            shapes.append(shape)

    result = AutoLabelingResult(shapes, replace=True)
    return result

def unload(self):
    del self.net

I recommand you to debug here.

[haitian-zhang]

大佬能说的具体点嘛，这部分我有个问题想问一下，我是直接复制了你示例unet.py的preprocess和predict_shapes部分，但是我自己的代码写的跟他不一样，需要换成我自己之前写的吗

[CVHub520]

Yes, you should use your own code if it's different. It's important to understand how your code works and ensure it's correct for your specific needs.