Detection object with custom YOLOv5 model by using SAHI: AttributeError: module 'yolov5' has no attribute 'load'

huangapple go评论96阅读模式
英文:

Detection object with custom YOLOv5 model by using SAHI: AttributeError: module 'yolov5' has no attribute 'load'

问题

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

  1. Traceback (most recent call last):
  2. File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py", line 29, in load_model
  3. model = yolov5.load(self.model_path, device=self.device)
  4. AttributeError: module 'yolov5' has no attribute 'load'
  5. During handling of the above exception, another exception occurred:
  6. Traceback (most recent call last):
  7. File "c:\Users\pawel\Documents\GitHub\AECVision\wall_detection_export_with_sahi.py", line 84, in <module>
  8. detection_model = AutoDetectionModel.from_pretrained(
  9. File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\auto_model.py", line 66, in from_pretrained
  10. return DetectionModel(
  11. File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\base.py", line 67, in __init__
  12. self.load_model()
  13. File "C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py", line 32, in load_model
  14. raise TypeError("model_path is not a valid yolov5 model path: ", e)
  15. TypeError: ('model_path is not a valid yolov5 model path: ', AttributeError("module 'yolov5' has no attribute 'load'"))

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

以下是我的代码:

  1. # 上传PDF并转换为JPG
  2. path_pdf = Path("wall_detection_export/upload_pdf")
  3. path_convert_pdf = Path("wall_detection_export/convert_pdf")
  4. path_export_txt = Path("wall_detection_export/export_txt")
  5. path_model = Path("train_results/model_12classes/weights/best.pt")
  6. converter = Convert_pdf(path_pdf=path_pdf)
  7. convert_file = converter.save_image(path_convert_pdf)
  8. # 设置检测模型
  9. detection_model = AutoDetectionModel.from_pretrained(
  10. model_type='yolov5',
  11. model_path=path_model,
  12. confidence_threshold=0.3,
  13. device="cuda", # 或者 'cuda:0'
  14. )
  15. # 使用SAHI切割预测
  16. result = get_sliced_prediction(
  17. convert_file,
  18. detection_model,
  19. slice_height = 1280,
  20. slice_width = 1280,
  21. overlap_height_ratio = 0.2,
  22. overlap_width_ratio = 0.2
  23. )
  24. result.export_visuals(export_dir=path_export_txt)

如何解决这个问题?感谢帮助!

英文:

I try to use SAHI library for object detection with my custom trained YOLOv5s6 model. I though SAHI support YOLOv5 models but when i try to build detection model i get an error:

  1. Traceback (most recent call last):
  2. File &quot;C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py&quot;, line 29, in load_model
  3. model = yolov5.load(self.model_path, device=self.device)
  4. AttributeError: module &#39;yolov5&#39; has no attribute &#39;load&#39;
  5. During handling of the above exception, another exception occurred:
  6. Traceback (most recent call last):
  7. File &quot;c:\Users\pawel\Documents\GitHub\AECVision\wall_detection_export_with_sahi.py&quot;, line 84, in &lt;module&gt;
  8. detection_model = AutoDetectionModel.from_pretrained(
  9. File &quot;C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\auto_model.py&quot;, line 66, in from_pretrained
  10. return DetectionModel(
  11. File &quot;C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\base.py&quot;, line 67, in __init__
  12. self.load_model()
  13. File &quot;C:\Users\pawel\Documents\GitHub\AECVision\aec-env\lib\site-packages\sahi\models\yolov5.py&quot;, line 32, in load_model
  14. raise TypeError(&quot;model_path is not a valid yolov5 model path: &quot;, e)
  15. TypeError: (&#39;model_path is not a valid yolov5 model path: &#39;, AttributeError(&quot;module &#39;yolov5&#39; has no attribute &#39;load&#39;&quot;))

I have my model weight in 'path_model'

Below is my code:

  1. # Upload pdf and change to jpg
  2. path_pdf = Path(&quot;wall_detection_export/upload_pdf&quot;)
  3. path_convert_pdf = Path(&quot;wall_detection_export/convert_pdf&quot;)
  4. path_export_txt = Path(&quot;wall_detection_export/export_txt&quot;)
  5. path_model = Path(&quot;train_results/model_12classes/weights/best.pt&quot;)
  6. converter = Convert_pdf(path_pdf=path_pdf)
  7. convert_file = converter.save_image(path_convert_pdf)
  8. # Set detection model
  9. detection_model = AutoDetectionModel.from_pretrained(
  10. model_type=&#39;yolov5&#39;,
  11. model_path=path_model,
  12. confidence_threshold=0.3,
  13. device=&quot;cuda&quot;, # or &#39;cuda:0&#39;
  14. )
  15. # Slice prediction with sahi
  16. result = get_sliced_prediction(
  17. convert_file,
  18. detection_model,
  19. slice_height = 1280,
  20. slice_width = 1280,
  21. overlap_height_ratio = 0.2,
  22. overlap_width_ratio = 0.2
  23. )
  24. result.export_visuals(export_dir=path_export_txt)

How can i fix this issue? Thanks for help!

答案1

得分: 0

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

yolov5_custom.py的代码如下:

  1. # OBSS SAHI Tool
  2. # Code written by Fatih C Akyon, 2020.
  3. import logging
  4. from typing import Any, Dict, List, Optional
  5. import numpy as np
  6. from sahi.models.base import DetectionModel
  7. from sahi.prediction import ObjectPrediction
  8. from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
  9. from sahi.utils.import_utils import check_package_minimum_version, check_requirements
  10. logger = logging.getLogger(__name__)
  11. class CustomYolov5DetectionModel(DetectionModel):
  12. def check_dependencies(self) -> None:
  13. check_requirements(["torch", "yolov5"])
  14. def load_model(self):
  15. """
  16. Detection model is initialized and set to self.model.
  17. """
  18. import torch
  19. try:
  20. model = torch.hub.load("yolov5", "custom", path=self.model_path, source="local")
  21. self.set_model(model)
  22. except Exception as e:
  23. raise TypeError("model_path is not a valid yolov5 model path: ", e)
  24. def set_model(self, model: Any):
  25. """
  26. Sets the underlying YOLOv5 model.
  27. Args:
  28. model: Any
  29. A YOLOv5 model
  30. """
  31. if model.__class__.__module__ not in ["yolov5.models.common", "models.common"]:
  32. raise Exception(f"Not a yolov5 model: {type(model)}")
  33. model.conf = self.confidence_threshold
  34. self.model = model
  35. # set category_mapping
  36. if not self.category_mapping:
  37. category_mapping = {str(ind): category_name for ind, category_name in enumerate(self.category_names)}
  38. self.category_mapping = category_mapping
  39. def perform_inference(self, image: np.ndarray):
  40. """
  41. Prediction is performed using self.model and the prediction result is set to self._original_predictions.
  42. Args:
  43. image: np.ndarray
  44. A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.
  45. """
  46. # Confirm model is loaded
  47. if self.model is None:
  48. raise ValueError("Model is not loaded, load it by calling .load_model()")
  49. if self.image_size is not None:
  50. prediction_result = self.model(image, size=self.image_size)
  51. else:
  52. prediction_result = self.model(image)
  53. self._original_predictions = prediction_result
  54. @property
  55. def num_categories(self):
  56. """
  57. Returns number of categories
  58. """
  59. return len(self.model.names)
  60. @property
  61. def has_mask(self):
  62. """
  63. Returns if model output contains segmentation mask
  64. """
  65. import yolov5
  66. from packaging import version
  67. if version.parse(yolov5.__version__) < version.parse("6.2.0"):
  68. return False
  69. else:
  70. return False # fix when yolov5 supports segmentation models
  71. @property
  72. def category_names(self):
  73. if check_package_minimum_version("yolov5", "6.2.0"):
  74. return list(self.model.names.values())
  75. else:
  76. return self.model.names
  77. def _create_object_prediction_list_from_original_predictions(
  78. self,
  79. shift_amount_list: Optional[List[List[int]]] = [[0, 0]],
  80. full_shape_list: Optional[List[List[int]]] = None,
  81. ):
  82. """
  83. self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
  84. self._object_prediction_list_per_image.
  85. Args:
  86. shift_amount_list: list of list
  87. To shift the box and mask predictions from sliced image to full sized image, should
  88. be in the form of List[[shift_x, shift_y],[shift_x, shift_y],...]
  89. full_shape_list: list of list
  90. Size of the full image after shifting, should be in the form of
  91. List[[height, width],[height, width],...]
  92. """
  93. original_predictions = self._original_predictions
  94. # compatilibty for sahi v0.8.15
  95. shift_amount_list = fix_shift_amount_list(shift_amount_list)
  96. full_shape_list = fix_full_shape_list(full_shape_list)
  97. # handle all predictions
  98. object_prediction_list_per_image = []
  99. for image_ind, image_predictions_in_xyxy_format in enumerate(original_predictions.xyxy):
  100. shift_amount = shift_amount_list[image_ind]
  101. full_shape = None if full_shape_list is None else full_shape_list[image_ind]
  102. object_prediction_list = []
  103. # process predictions
  104. for prediction in image_predictions_in_xyxy_format.cpu().detach().numpy():
  105. x1 = prediction[0]
  106. y1 = prediction[1]
  107. x2 = prediction[2]
  108. y2 = prediction[3]
  109. bbox = [x1, y1, x2, y2]
  110. score = prediction[4]
  111. category_id = int(prediction[5])
  112. category_name = self.category_mapping[str(category_id)]
  113. # fix negative box coords
  114. bbox[0] = max(0, bbox[0])
  115. bbox[1] = max(0, bbox[1])
  116. bbox[2] = max(0, bbox[2])
  117. bbox[3] = max(0, bbox[3])
  118. # fix out of image box coords
  119. if full_shape is not None:
  120. bbox[0] = min(full_shape[1], bbox[0])
  121. bbox[1] = min(full_shape[0], bbox[1])
  122. bbox[2] = min(full_shape[1], bbox[2])
  123. bbox[3] = min(full_shape[0], bbox[3])
  124. # ignore invalid predictions
  125. if not (bbox[0] < bbox[2]) or not (bbox[1] < bbox[3]):
  126. logger.warning(f"ignoring invalid prediction with bbox: {bbox}")
  127. continue
  128. object_prediction = ObjectPrediction(
  129. bbox=bbox,
  130. category_id=category_id,
  131. score=score,
  132. bool_mask=None,
  133. category_name=category_name,
  134. shift_amount=shift_amount,
  135. full_shape=full_shape,
  136. )
  137. object_prediction_list.append(object_prediction)
  138. object_prediction_list_per_image.append(object_prediction_list)
  139. self._object_prediction_list_per_image = object_prediction_list_per_image

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

  1. MODEL_TYPE_TO_MODEL_CLASS_NAME = {
  2. "yolov8": "Yolov8DetectionModel",
  3. "mmdet": "MmdetDetectionModel",
  4. "yolov5": "Yolov5DetectionModel",
  5. "detectron2": "Detectron2DetectionModel",
  6. "huggingface": "HuggingfaceDetectionModel",
  7. "torchvision": "TorchVisionDetectionModel",
  8. "yolov5sparse": "Yolov5SparseDetectionModel",
  9. "yolonas": "YoloNasDetectionModel",
  10. "yolov5_custom": "CustomYolov5DetectionModel"
  11. }

请确保

英文:

You need to add two thing in Sahi library in your environment: yolov5_custom.py (class with your model) and add your model to dictionary in auto_model.py

Detection object with custom YOLOv5 model by using SAHI: AttributeError: module 'yolov5' has no attribute 'load'

Below i place code in:
yolov5_custom.py

  1. # OBSS SAHI Tool
  2. # Code written by Fatih C Akyon, 2020.
  3. import logging
  4. from typing import Any, Dict, List, Optional
  5. import numpy as np
  6. from sahi.models.base import DetectionModel
  7. from sahi.prediction import ObjectPrediction
  8. from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
  9. from sahi.utils.import_utils import check_package_minimum_version, check_requirements
  10. logger = logging.getLogger(__name__)
  11. class CustomYolov5DetectionModel(DetectionModel):
  12. def check_dependencies(self) -&gt; None:
  13. check_requirements([&quot;torch&quot;, &quot;yolov5&quot;])
  14. def load_model(self):
  15. &quot;&quot;&quot;
  16. Detection model is initialized and set to self.model.
  17. &quot;&quot;&quot;
  18. import torch
  19. try:
  20. model = torch.hub.load(&quot;yolov5&quot;, &quot;custom&quot;, path=self.model_path, source=&quot;local&quot;)
  21. self.set_model(model)
  22. except Exception as e:
  23. raise TypeError(&quot;model_path is not a valid yolov5 model path: &quot;, e)
  24. def set_model(self, model: Any):
  25. &quot;&quot;&quot;
  26. Sets the underlying YOLOv5 model.
  27. Args:
  28. model: Any
  29. A YOLOv5 model
  30. &quot;&quot;&quot;
  31. if model.__class__.__module__ not in [&quot;yolov5.models.common&quot;, &quot;models.common&quot;]:
  32. raise Exception(f&quot;Not a yolov5 model: {type(model)}&quot;)
  33. model.conf = self.confidence_threshold
  34. self.model = model
  35. # set category_mapping
  36. if not self.category_mapping:
  37. category_mapping = {str(ind): category_name for ind, category_name in enumerate(self.category_names)}
  38. self.category_mapping = category_mapping
  39. def perform_inference(self, image: np.ndarray):
  40. &quot;&quot;&quot;
  41. Prediction is performed using self.model and the prediction result is set to self._original_predictions.
  42. Args:
  43. image: np.ndarray
  44. A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.
  45. &quot;&quot;&quot;
  46. # Confirm model is loaded
  47. if self.model is None:
  48. raise ValueError(&quot;Model is not loaded, load it by calling .load_model()&quot;)
  49. if self.image_size is not None:
  50. prediction_result = self.model(image, size=self.image_size)
  51. else:
  52. prediction_result = self.model(image)
  53. self._original_predictions = prediction_result
  54. @property
  55. def num_categories(self):
  56. &quot;&quot;&quot;
  57. Returns number of categories
  58. &quot;&quot;&quot;
  59. return len(self.model.names)
  60. @property
  61. def has_mask(self):
  62. &quot;&quot;&quot;
  63. Returns if model output contains segmentation mask
  64. &quot;&quot;&quot;
  65. import yolov5
  66. from packaging import version
  67. if version.parse(yolov5.__version__) &lt; version.parse(&quot;6.2.0&quot;):
  68. return False
  69. else:
  70. return False # fix when yolov5 supports segmentation models
  71. @property
  72. def category_names(self):
  73. if check_package_minimum_version(&quot;yolov5&quot;, &quot;6.2.0&quot;):
  74. return list(self.model.names.values())
  75. else:
  76. return self.model.names
  77. def _create_object_prediction_list_from_original_predictions(
  78. self,
  79. shift_amount_list: Optional[List[List[int]]] = [[0, 0]],
  80. full_shape_list: Optional[List[List[int]]] = None,
  81. ):
  82. &quot;&quot;&quot;
  83. self._original_predictions is converted to a list of prediction.ObjectPrediction and set to
  84. self._object_prediction_list_per_image.
  85. Args:
  86. shift_amount_list: list of list
  87. To shift the box and mask predictions from sliced image to full sized image, should
  88. be in the form of List[[shift_x, shift_y],[shift_x, shift_y],...]
  89. full_shape_list: list of list
  90. Size of the full image after shifting, should be in the form of
  91. List[[height, width],[height, width],...]
  92. &quot;&quot;&quot;
  93. original_predictions = self._original_predictions
  94. # compatilibty for sahi v0.8.15
  95. shift_amount_list = fix_shift_amount_list(shift_amount_list)
  96. full_shape_list = fix_full_shape_list(full_shape_list)
  97. # handle all predictions
  98. object_prediction_list_per_image = []
  99. for image_ind, image_predictions_in_xyxy_format in enumerate(original_predictions.xyxy):
  100. shift_amount = shift_amount_list[image_ind]
  101. full_shape = None if full_shape_list is None else full_shape_list[image_ind]
  102. object_prediction_list = []
  103. # process predictions
  104. for prediction in image_predictions_in_xyxy_format.cpu().detach().numpy():
  105. x1 = prediction[0]
  106. y1 = prediction[1]
  107. x2 = prediction[2]
  108. y2 = prediction[3]
  109. bbox = [x1, y1, x2, y2]
  110. score = prediction[4]
  111. category_id = int(prediction[5])
  112. category_name = self.category_mapping[str(category_id)]
  113. # fix negative box coords
  114. bbox[0] = max(0, bbox[0])
  115. bbox[1] = max(0, bbox[1])
  116. bbox[2] = max(0, bbox[2])
  117. bbox[3] = max(0, bbox[3])
  118. # fix out of image box coords
  119. if full_shape is not None:
  120. bbox[0] = min(full_shape[1], bbox[0])
  121. bbox[1] = min(full_shape[0], bbox[1])
  122. bbox[2] = min(full_shape[1], bbox[2])
  123. bbox[3] = min(full_shape[0], bbox[3])
  124. # ignore invalid predictions
  125. if not (bbox[0] &lt; bbox[2]) or not (bbox[1] &lt; bbox[3]):
  126. logger.warning(f&quot;ignoring invalid prediction with bbox: {bbox}&quot;)
  127. continue
  128. object_prediction = ObjectPrediction(
  129. bbox=bbox,
  130. category_id=category_id,
  131. score=score,
  132. bool_mask=None,
  133. category_name=category_name,
  134. shift_amount=shift_amount,
  135. full_shape=full_shape,
  136. )
  137. object_prediction_list.append(object_prediction)
  138. object_prediction_list_per_image.append(object_prediction_list)
  139. self._object_prediction_list_per_image = object_prediction_list_per_image

And add your new model to dict in auto_model.py

  1. MODEL_TYPE_TO_MODEL_CLASS_NAME = {
  2. &quot;yolov8&quot;: &quot;Yolov8DetectionModel&quot;,
  3. &quot;mmdet&quot;: &quot;MmdetDetectionModel&quot;,
  4. &quot;yolov5&quot;: &quot;Yolov5DetectionModel&quot;,
  5. &quot;detectron2&quot;: &quot;Detectron2DetectionModel&quot;,
  6. &quot;huggingface&quot;: &quot;HuggingfaceDetectionModel&quot;,
  7. &quot;torchvision&quot;: &quot;TorchVisionDetectionModel&quot;,
  8. &quot;yolov5sparse&quot;: &quot;Yolov5SparseDetectionModel&quot;,
  9. &quot;yolonas&quot;: &quot;YoloNasDetectionModel&quot;,
  10. &quot;yolov5_custom&quot;: &quot;CustomYolov5DetectionModel&quot;
  11. }

huangapple
  • 本文由 发表于 2023年7月23日 19:43:37
  • 转载请务必保留本文链接:https://go.coder-hub.com/76748038.html
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定