Sahi .pt and .pth problem，can not predict

[nakuda1]

Dear Sahi contributor, recently I train my own yolov5 model and got .pth to do predict. So far, I rename .pth to .pt and use sahi to do sliceing. But it does not work. The error is below. Can you help me to explain what can I do to use Sahi.

KeyError Traceback (most recent call last)
File ~/miniconda3/lib/python3.8/site-packages/sahi/model.py:370, in Yolov5DetectionModel.load_model(self)
369 try:
--> 370 model = yolov5.load(self.model_path, device=self.device)
371 self.set_model(model)

File ~/miniconda3/lib/python3.8/site-packages/yolov5/helpers.py:34, in load_model(model_path, device, autoshape, verbose)
32 device = torch.device(device)
---> 34 model = DetectMultiBackend(model_path, device=device)
36 if autoshape:

File ~/miniconda3/lib/python3.8/site-packages/yolov5/models/common.py:310, in DetectMultiBackend.init(self, weights, device, dnn, data, fp16)
309 if pt: # PyTorch
--> 310 model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
311 stride = max(int(model.stride.max()), 32) # model stride

File ~/miniconda3/lib/python3.8/site-packages/yolov5/models/experimental.py:100, in attempt_load(weights, map_location, inplace, fuse)
99 ckpt = torch.load(attempt_download(w), map_location=map_location) # load
--> 100 ckpt = (ckpt.get('ema') or ckpt['model']).float() # FP32 model
101 model.append(ckpt.fuse().eval() if fuse else ckpt.eval()) # fused or un-fused model in eval mode

KeyError: 'model'

During handling of the above exception, another exception occurred:

TypeError Traceback (most recent call last)
Input In [2], in <cell line: 5>()
2 from sahi.predict import get_sliced_prediction
4 # 初始化检测模型
----> 5 detection_model = AutoDetectionModel.from_pretrained(
6 model_type='yolov5',
7 model_path='best_epoch_weights.pt',
8 confidence_threshold=0.3,
9 device="cpu", # or 'cuda:0'
10 )
11 image = './1.jpg'
14 result = get_sliced_prediction(
15 image,
16 detection_model,
(...)
23 postprocess_class_agnostic = True,
24 )

File ~/miniconda3/lib/python3.8/site-packages/sahi/auto_model.py:63, in AutoDetectionModel.from_pretrained(model_type, model_path, config_path, device, mask_threshold, confidence_threshold, category_mapping, category_remapping, load_at_init, image_size, **kwargs)
60 model_class_name = MODEL_TYPE_TO_MODEL_CLASS_NAME[model_type]
61 DetectionModel = import_model_class(model_class_name)
---> 63 return DetectionModel(
64 model_path=model_path,
65 config_path=config_path,
66 device=device,
67 mask_threshold=mask_threshold,
68 confidence_threshold=confidence_threshold,
69 category_mapping=category_mapping,
70 category_remapping=category_remapping,
71 load_at_init=load_at_init,
72 image_size=image_size,
73 **kwargs,
74 )

File ~/miniconda3/lib/python3.8/site-packages/sahi/model.py:77, in DetectionModel.init(self, model_path, model, config_path, device, mask_threshold, confidence_threshold, category_mapping, category_remapping, load_at_init, image_size)
75 self.set_model(model)
76 else:
---> 77 self.load_model()

File ~/miniconda3/lib/python3.8/site-packages/sahi/model.py:373, in Yolov5DetectionModel.load_model(self)
371 self.set_model(model)
372 except Exception as e:
--> 373 raise TypeError("model_path is not a valid yolov5 model path: ", e)

TypeError: ('model_path is not a valid yolov5 model path: ', KeyError('model'))

[809441282]

Do you figure out this problem?l meet the seem one. sometimes, the detection can run correctly in the form of pt.

[nakuda1]

Yes, I figure it. I use a new pth and write the custormrize model with sahi myself.