Spaces:
Running
Running
Delete omniparser.py
Browse files- 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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|