This commit is contained in:
yadonglu
2024-10-09 22:57:41 +00:00
parent 741af789a8
commit 5df0bc8005
2 changed files with 4 additions and 4 deletions

View File

@@ -385,7 +385,7 @@
"from PIL import Image\n", "from PIL import Image\n",
"device = 'cuda'\n", "device = 'cuda'\n",
"\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" "som_model.to(device)\n"
] ]
}, },
@@ -404,7 +404,7 @@
], ],
"source": [ "source": [
"\n", "\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",
"\n" "\n"
] ]

View File

@@ -12,8 +12,8 @@ from utils import check_ocr_box, get_yolo_model, get_caption_model_processor, ge
import torch import torch
from PIL import Image from PIL import Image
yolo_model = get_yolo_model(model_path='weights/omniparser/weights/best.pt') 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/blipv2_ui_merge", device='cuda') caption_model_processor = get_caption_model_processor(model_name_or_path="weights/omniparser/icon_caption_blip2", device='cuda')
platform = 'pc' platform = 'pc'
if platform == 'pc': if platform == 'pc':
draw_bbox_config = { draw_bbox_config = {