update eadme, safetensor

.gitignore

@@ -1,4 +1,5 @@
 weights/icon_caption_blip2
 weights/icon_caption_florence
+weights/icon_detect
 .gradio
 __pycache__

README.md

@@ -16,13 +16,20 @@
 - [2024/09] OmniParser achieves the best performance on [Windows Agent Arena](https://microsoft.github.io/WindowsAgentArena/)! 
 
 ## Install 
+Install environment:
 ```python
 conda create -n "omni" python==3.12
 conda activate omni
 pip install -r requirement.txt
 ```
+
 Then download the model ckpts files in: https://huggingface.co/microsoft/OmniParser, and put them under weights/, default folder structure is: weights/icon_detect, weights/icon_caption_florence, weights/icon_caption_blip2. 
 
+Finally, convert the safetensor to .pt file. 
+```python
+python weights/convert_safetensor_to_pt.py
+```
+
 ## Examples:
 We put together a few simple examples in the demo.ipynb. 
 

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')