avans06 commited on
Commit
4e136fc
·
1 Parent(s): 5d975b7

fix: allow loading YOLOv5 custom model by adding DetectionModel to torch safe globals

Browse files

PyTorch 2.6+ enforces stricter pickling security, which prevents loading models containing custom classes like `DetectionModel` unless they are explicitly allowlisted.

Files changed (1) hide show
  1. image_processing/model.py +3 -0
image_processing/model.py CHANGED
@@ -11,10 +11,13 @@ class Model:
11
  if not self.imported:
12
  self.imported = True
13
  import torch
 
14
  import pathlib
15
  import sys
16
  import os
17
  from myutils.respath import resource_path
 
 
18
 
19
  # Redirect sys.stderr to a file or a valid stream
20
  if sys.stderr is None:
 
11
  if not self.imported:
12
  self.imported = True
13
  import torch
14
+ import torch.serialization
15
  import pathlib
16
  import sys
17
  import os
18
  from myutils.respath import resource_path
19
+ from yolov5.models.yolo import DetectionModel
20
+ torch.serialization.add_safe_globals([DetectionModel]) # Exception: Weights only load failed.
21
 
22
  # Redirect sys.stderr to a file or a valid stream
23
  if sys.stderr is None: