gaoyu1314 commited on
Commit
a81d82e
·
verified ·
1 Parent(s): 7f4f003

Delete omniparser.py

Browse files
Files changed (1) hide show
  1. omniparser.py +0 -32
omniparser.py DELETED
@@ -1,32 +0,0 @@
1
- from util.utils import get_som_labeled_img, get_caption_model_processor, get_yolo_model, check_ocr_box
2
- import torch
3
- from PIL import Image
4
- import io
5
- import base64
6
- from typing import Dict
7
- class Omniparser(object):
8
- def __init__(self, config: Dict):
9
- self.config = config
10
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
11
-
12
- self.som_model = get_yolo_model(model_path=config['som_model_path'])
13
- self.caption_model_processor = get_caption_model_processor(model_name=config['caption_model_name'], model_name_or_path=config['caption_model_path'], device=device)
14
- print('Omniparser initialized!!!')
15
-
16
- def parse(self, image_base64: str):
17
- image_bytes = base64.b64decode(image_base64)
18
- image = Image.open(io.BytesIO(image_bytes))
19
- print('image size:', image.size)
20
-
21
- box_overlay_ratio = max(image.size) / 3200
22
- draw_bbox_config = {
23
- 'text_scale': 0.8 * box_overlay_ratio,
24
- 'text_thickness': max(int(2 * box_overlay_ratio), 1),
25
- 'text_padding': max(int(3 * box_overlay_ratio), 1),
26
- 'thickness': max(int(3 * box_overlay_ratio), 1),
27
- }
28
-
29
- (text, ocr_bbox), _ = check_ocr_box(image, display_img=False, output_bb_format='xyxy', easyocr_args={'text_threshold': 0.8}, use_paddleocr=False)
30
- dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image, self.som_model, BOX_TRESHOLD = self.config['BOX_TRESHOLD'], output_coord_in_ratio=True, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=self.caption_model_processor, ocr_text=text,use_local_semantics=True, iou_threshold=0.7, scale_img=False, batch_size=128)
31
-
32
- return dino_labled_img, parsed_content_list