问题

我尝试使用SAHI库来进行目标检测，使用了我自己训练的YOLOv5s6模型。我认为SAHI支持YOLOv5模型，但当我尝试构建检测模型时出现了错误：

Traceback (most recent call last):
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py", line 29, in load_model
    model = yolov5.load(self.model_path, device=self.device)
AttributeError: module 'yolov5' has no attribute 'load'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "c:\Users\pawel\Documents\GitHub\AECVision\wall_detection_export_with_sahi.py", line 84, in <module>
    detection_model = AutoDetectionModel.from_pretrained(
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\auto_model.py", line 66, in from_pretrained
    return DetectionModel(
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\base.py", line 67, in __init__
    self.load_model()
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py", line 32, in load_model
    raise TypeError("model_path is not a valid yolov5 model path: ", e)
TypeError: ('model_path is not a valid yolov5 model path: ', AttributeError("module 'yolov5' has no attribute 'load'"))

我将我的模型权重保存在'path_model'中。

以下是我的代码：

# 上传PDF并转换为JPG
path_pdf = Path("wall_detection_export/upload_pdf")
path_convert_pdf = Path("wall_detection_export/convert_pdf")
path_export_txt = Path("wall_detection_export/export_txt")
path_model = Path("train_results/model_12classes/weights/best.pt")

converter = Convert_pdf(path_pdf=path_pdf)
convert_file = converter.save_image(path_convert_pdf)

# 设置检测模型
detection_model = AutoDetectionModel.from_pretrained(
    model_type='yolov5',
    model_path=path_model,
    confidence_threshold=0.3,
    device="cuda", # 或者 'cuda:0'
)

# 使用SAHI切割预测
result = get_sliced_prediction(
    convert_file,
    detection_model,
    slice_height = 1280,
    slice_width = 1280,
    overlap_height_ratio = 0.2,
    overlap_width_ratio = 0.2
)
result.export_visuals(export_dir=path_export_txt)

如何解决这个问题？感谢帮助！

Traceback (most recent call last):
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py", line 29, in load_model
    model = yolov5.load(self.model_path, device=self.device)
AttributeError: module 'yolov5' has no attribute 'load'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "c:\Users\pawel\Documents\GitHub\AECVision\wall_detection_export_with_sahi.py", line 84, in <module>
    detection_model = AutoDetectionModel.from_pretrained(
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\auto_model.py", line 66, in from_pretrained
    return DetectionModel(
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\base.py", line 67, in __init__
    self.load_model()
  File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py", line 32, in load_model
    raise TypeError("model_path is not a valid yolov5 model path: ", e)
TypeError: ('model_path is not a valid yolov5 model path: ', AttributeError("module 'yolov5' has no attribute 'load'"))

# Upload pdf and change to jpg
path_pdf = Path("wall_detection_export/upload_pdf")
path_convert_pdf = Path("wall_detection_export/convert_pdf")
path_export_txt = Path("wall_detection_export/export_txt")
path_model = Path("train_results/model_12classes/weights/best.pt")

converter = Convert_pdf(path_pdf=path_pdf)
convert_file = converter.save_image(path_convert_pdf)

# Set detection model
detection_model = AutoDetectionModel.from_pretrained(
    model_type='yolov5',
    model_path=path_model,
    confidence_threshold=0.3,
    device="cuda", # or 'cuda:0'
)

# Slice prediction with sahi
result = get_sliced_prediction(
    convert_file,
    detection_model,
    slice_height = 1280,
    slice_width = 1280,
    overlap_height_ratio = 0.2,
    overlap_width_ratio = 0.2
)
result.export_visuals(export_dir=path_export_txt)

答案1

得分: 0

你需要在Sahi库中的环境中添加两个东西：yolov5_custom.py（包含您的模型的类）和将您的模型添加到auto_model.py中的字典中。

yolov5_custom.py的代码如下：

# OBSS SAHI Tool
# Code written by Fatih C Akyon, 2020.

import logging
from typing import Any, Dict, List, Optional

import numpy as np

from sahi.models.base import DetectionModel
from sahi.prediction import ObjectPrediction
from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
from sahi.utils.import_utils import check_package_minimum_version, check_requirements

logger = logging.getLogger(__name__)


class CustomYolov5DetectionModel(DetectionModel):
    def check_dependencies(self) -> None:
        check_requirements(["torch", "yolov5"])

    def load_model(self):
        """
        Detection model is initialized and set to self.model.
        """

        import torch

        try:
            model = torch.hub.load("yolov5", "custom", path=self.model_path, source="local")
            self.set_model(model)
        except Exception as e:
            raise TypeError("model_path is not a valid yolov5 model path: ", e)

    def set_model(self, model: Any):
        """
        Sets the underlying YOLOv5 model.
        Args:
            model: Any
                A YOLOv5 model
        """

        if model.__class__.__module__ not in ["yolov5.models.common", "models.common"]:
            raise Exception(f"Not a yolov5 model: {type(model)}")

        model.conf = self.confidence_threshold
        self.model = model

        # set category_mapping
        if not self.category_mapping:
            category_mapping = {str(ind): category_name for ind, category_name in enumerate(self.category_names)}
            self.category_mapping = category_mapping

    def perform_inference(self, image: np.ndarray):
        """
        Prediction is performed using self.model and the prediction result is set to self._original_predictions.
        Args:
            image: np.ndarray
                A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.
        """

        # Confirm model is loaded
        if self.model is None:
            raise ValueError("Model is not loaded, load it by calling .load_model()")
        if self.image_size is not None:
            prediction_result = self.model(image, size=self.image_size)
        else:
            prediction_result = self.model(image)

        self._original_predictions = prediction_result

    @property
    def num_categories(self):
        """
        Returns number of categories
        """
        return len(self.model.names)

    @property
    def has_mask(self):
        """
        Returns if model output contains segmentation mask
        """
        import yolov5
        from packaging import version

        if version.parse(yolov5.__version__) < version.parse("6.2.0"):
            return False
        else:
            return False  # fix when yolov5 supports segmentation models

    @property
    def category_names(self):
        if check_package_minimum_version("yolov5", "6.2.0"):
            return list(self.model.names.values())
        else:
            return self.model.names

    def _create_object_prediction_list_from_original_predictions(
        self,
        shift_amount_list: Optional[List[List[int]]] = [[0, 0]],
        full_shape_list: Optional[List[List[int]]] = None,
    ):
        """
        self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
        self._object_prediction_list_per_image.
        Args:
            shift_amount_list: list of list
                To shift the box and mask predictions from sliced image to full sized image, should
                be in the form of List[[shift_x, shift_y],[shift_x, shift_y],...]
            full_shape_list: list of list
                Size of the full image after shifting, should be in the form of
                List[[height, width],[height, width],...]
        """
        original_predictions = self._original_predictions

        # compatilibty for sahi v0.8.15
        shift_amount_list = fix_shift_amount_list(shift_amount_list)
        full_shape_list = fix_full_shape_list(full_shape_list)

        # handle all predictions
        object_prediction_list_per_image = []
        for image_ind, image_predictions_in_xyxy_format in enumerate(original_predictions.xyxy):
            shift_amount = shift_amount_list[image_ind]
            full_shape = None if full_shape_list is None else full_shape_list[image_ind]
            object_prediction_list = []

            # process predictions
            for prediction in image_predictions_in_xyxy_format.cpu().detach().numpy():
                x1 = prediction[0]
                y1 = prediction[1]
                x2 = prediction[2]
                y2 = prediction[3]
                bbox = [x1, y1, x2, y2]
                score = prediction[4]
                category_id = int(prediction[5])
                category_name = self.category_mapping[str(category_id)]

                # fix negative box coords
                bbox[0] = max(0, bbox[0])
                bbox[1] = max(0, bbox[1])
                bbox[2] = max(0, bbox[2])
                bbox[3] = max(0, bbox[3])

                # fix out of image box coords
                if full_shape is not None:
                    bbox[0] = min(full_shape[1], bbox[0])
                    bbox[1] = min(full_shape[0], bbox[1])
                    bbox[2] = min(full_shape[1], bbox[2])
                    bbox[3] = min(full_shape[0], bbox[3])

                # ignore invalid predictions
                if not (bbox[0] < bbox[2]) or not (bbox[1] < bbox[3]):
                    logger.warning(f"ignoring invalid prediction with bbox: {bbox}")
                    continue

                object_prediction = ObjectPrediction(
                    bbox=bbox,
                    category_id=category_id,
                    score=score,
                    bool_mask=None,
                    category_name=category_name,
                    shift_amount=shift_amount,
                    full_shape=full_shape,
                )
                object_prediction_list.append(object_prediction)
            object_prediction_list_per_image.append(object_prediction_list)

        self._object_prediction_list_per_image = object_prediction_list_per_image

然后，将您的新模型添加到auto_model.py中的字典中：

MODEL_TYPE_TO_MODEL_CLASS_NAME = {
    "yolov8": "Yolov8DetectionModel",
    "mmdet": "MmdetDetectionModel",
    "yolov5": "Yolov5DetectionModel",
    "detectron2": "Detectron2DetectionModel",
    "huggingface": "HuggingfaceDetectionModel",
    "torchvision": "TorchVisionDetectionModel",
    "yolov5sparse": "Yolov5SparseDetectionModel",
    "yolonas": "YoloNasDetectionModel",
    "yolov5_custom": "CustomYolov5DetectionModel"
}

请确保

# OBSS SAHI Tool
# Code written by Fatih C Akyon, 2020.
import logging
from typing import Any, Dict, List, Optional
import numpy as np
from sahi.models.base import DetectionModel
from sahi.prediction import ObjectPrediction
from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
from sahi.utils.import_utils import check_package_minimum_version, check_requirements
logger = logging.getLogger(__name__)
class CustomYolov5DetectionModel(DetectionModel):
def check_dependencies(self) -> None:
check_requirements(["torch", "yolov5"])
def load_model(self):
"""
Detection model is initialized and set to self.model.
"""
import torch
try:
model = torch.hub.load("yolov5", "custom", path=self.model_path, source="local")
self.set_model(model)
except Exception as e:
raise TypeError("model_path is not a valid yolov5 model path: ", e)
def set_model(self, model: Any):
"""
Sets the underlying YOLOv5 model.
Args:
model: Any
A YOLOv5 model
"""
if model.__class__.__module__ not in ["yolov5.models.common", "models.common"]:
raise Exception(f"Not a yolov5 model: {type(model)}")
model.conf = self.confidence_threshold
self.model = model
# set category_mapping
if not self.category_mapping:
category_mapping = {str(ind): category_name for ind, category_name in enumerate(self.category_names)}
self.category_mapping = category_mapping
def perform_inference(self, image: np.ndarray):
"""
Prediction is performed using self.model and the prediction result is set to self._original_predictions.
Args:
image: np.ndarray
A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.
"""
# Confirm model is loaded
if self.model is None:
raise ValueError("Model is not loaded, load it by calling .load_model()")
if self.image_size is not None:
prediction_result = self.model(image, size=self.image_size)
else:
prediction_result = self.model(image)
self._original_predictions = prediction_result
@property
def num_categories(self):
"""
Returns number of categories
"""
return len(self.model.names)
@property
def has_mask(self):
"""
Returns if model output contains segmentation mask
"""
import yolov5
from packaging import version
if version.parse(yolov5.__version__) < version.parse("6.2.0"):
return False
else:
return False  # fix when yolov5 supports segmentation models
@property
def category_names(self):
if check_package_minimum_version("yolov5", "6.2.0"):
return list(self.model.names.values())
else:
return self.model.names
def _create_object_prediction_list_from_original_predictions(
self,
shift_amount_list: Optional[List[List[int]]] = [[0, 0]],
full_shape_list: Optional[List[List[int]]] = None,
):
"""
self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
self._object_prediction_list_per_image.
Args:
shift_amount_list: list of list
To shift the box and mask predictions from sliced image to full sized image, should
be in the form of List[[shift_x, shift_y],[shift_x, shift_y],...]
full_shape_list: list of list
Size of the full image after shifting, should be in the form of
List[[height, width],[height, width],...]
"""
original_predictions = self._original_predictions
# compatilibty for sahi v0.8.15
shift_amount_list = fix_shift_amount_list(shift_amount_list)
full_shape_list = fix_full_shape_list(full_shape_list)
# handle all predictions
object_prediction_list_per_image = []
for image_ind, image_predictions_in_xyxy_format in enumerate(original_predictions.xyxy):
shift_amount = shift_amount_list[image_ind]
full_shape = None if full_shape_list is None else full_shape_list[image_ind]
object_prediction_list = []
# process predictions
for prediction in image_predictions_in_xyxy_format.cpu().detach().numpy():
x1 = prediction[0]
y1 = prediction[1]
x2 = prediction[2]
y2 = prediction[3]
bbox = [x1, y1, x2, y2]
score = prediction[4]
category_id = int(prediction[5])
category_name = self.category_mapping[str(category_id)]
# fix negative box coords
bbox[0] = max(0, bbox[0])
bbox[1] = max(0, bbox[1])
bbox[2] = max(0, bbox[2])
bbox[3] = max(0, bbox[3])
# fix out of image box coords
if full_shape is not None:
bbox[0] = min(full_shape[1], bbox[0])
bbox[1] = min(full_shape[0], bbox[1])
bbox[2] = min(full_shape[1], bbox[2])
bbox[3] = min(full_shape[0], bbox[3])
# ignore invalid predictions
if not (bbox[0] < bbox[2]) or not (bbox[1] < bbox[3]):
logger.warning(f"ignoring invalid prediction with bbox: {bbox}")
continue
object_prediction = ObjectPrediction(
bbox=bbox,
category_id=category_id,
score=score,
bool_mask=None,
category_name=category_name,
shift_amount=shift_amount,
full_shape=full_shape,
)
object_prediction_list.append(object_prediction)
object_prediction_list_per_image.append(object_prediction_list)
self._object_prediction_list_per_image = object_prediction_list_per_image

MODEL_TYPE_TO_MODEL_CLASS_NAME = {
"yolov8": "Yolov8DetectionModel",
"mmdet": "MmdetDetectionModel",
"yolov5": "Yolov5DetectionModel",
"detectron2": "Detectron2DetectionModel",
"huggingface": "HuggingfaceDetectionModel",
"torchvision": "TorchVisionDetectionModel",
"yolov5sparse": "Yolov5SparseDetectionModel",
"yolonas": "YoloNasDetectionModel",
"yolov5_custom": "CustomYolov5DetectionModel"
}