minor:

yanlianfu · Oct 9, 2024 · 5df0bc8 · 5df0bc8
1 parent 741af78
commit 5df0bc8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/demo.ipynb b/demo.ipynb
@@ -385,7 +385,7 @@
     "from PIL import Image\n",
     "device = 'cuda'\n",
     "\n",
-    "som_model = get_yolo_model(model_path='weights/omniparser/weights/best.pt')\n",
+    "som_model = get_yolo_model(model_path='weights/omniparser/icon_detect/best.pt')\n",
     "som_model.to(device)\n"
    ]
   },
@@ -404,7 +404,7 @@
    ],
    "source": [
     "\n",
-    "caption_model_processor = get_caption_model_processor(model_name_or_path=\"weights/omniparser/blipv2_ui_merge\", device=device)\n",
+    "caption_model_processor = get_caption_model_processor(model_name_or_path=\"weights/omniparser/icon_caption_blip2\", device=device)\n",
     "\n",
     "\n"
    ]

diff --git a/gradio_demo.py b/gradio_demo.py
@@ -12,8 +12,8 @@
 import torch
 from PIL import Image
 
-yolo_model = get_yolo_model(model_path='weights/omniparser/weights/best.pt')
-caption_model_processor = get_caption_model_processor(model_name_or_path="weights/omniparser/blipv2_ui_merge", device='cuda')
+yolo_model = get_yolo_model(model_path='weights/omniparser/icon_caption_blip2/best.pt')
+caption_model_processor = get_caption_model_processor(model_name_or_path="weights/omniparser/icon_caption_blip2", device='cuda')
 platform = 'pc'
 if platform == 'pc':
     draw_bbox_config = {