update eadme, safetensor

2024-10-26 14:09:29 -07:00
parent 65e14323d1
commit 8a3c07d0cd
5 changed files with 75 additions and 419 deletions
--- a/weights/convert_safetensor_to_pt.py
+++ b/weights/convert_safetensor_to_pt.py
@@ -0,0 +1,9 @@
+import torch
+from ultralytics.nn.tasks import DetectionModel
+from safetensors.torch import load_file
+
+tensor_dict = load_file("weights/icon_detect/model.safetensors")
+
+model = DetectionModel('weights/icon_detect/model.yaml')
+model.load_state_dict(tensor_dict)
+torch.save({'model':model}, 'weights/icon_detect/best.pt')
--- a/weights/icon_detect/best.pt
+++ b/weights/icon_detect/best.pt