Dom/OmniParser
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update readme

Browse Source
This commit is contained in:
yadonglu
2024-12-09 12:26:04 -08:00
parent 075f349ea1
commit 669efd3611
6 changed files with 1149 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
weights/convert_safetensor_to_pt.py
View File

@@ -5,16 +5,13 @@ import argparse
import yaml
import os
# accept args to specify v1 or v1_5
parser = argparse.ArgumentParser(description='Specify version v1 or v1_5')
parser.add_argument('--weights_dir', type=str, required=True, help='Specify the path to the safetensor file', default='weights/icon_detect_v1_5')
# accept args to specify v1
parser = argparse.ArgumentParser(description='add weight directory')
parser.add_argument('--weights_dir', type=str, required=True, help='Specify the path to the safetensor file', default='weights/icon_detect')
args = parser.parse_args()
tensor_dict = load_file(os.path.join(args.weights_dir, "model.safetensors"))
model = DetectionModel(os.path.join(args.weights_dir, "model.yaml"))
# from ultralytics import YOLO
# som_model = YOLO("yolo11m.pt")
# model = som_model.model
model.load_state_dict(tensor_dict)
save_dict = {'model':model}