我尝试使用SAHI库对我的自定义训练YOLOv 5s 6模型进行对象检测。我虽然SAHI支持YOLOv 5模型,但当我尝试构建检测模型时,我得到一个错误:
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'中
下面是我的代码:
# 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条答案
按热度按时间xmjla07d1#
您需要在您的环境中的Sahi库中添加两个东西:yolov5_custom.py(使用您的模型创建类),并将您的模型添加到auto_model.py中的字典中
的数据
下面我把代码放在:yolov5_custom.py
字符串
并将新模型添加到auto_model.py中的dict中
型