[email protected] commited on
Commit
3381ab5
·
0 Parent(s):

Added model files

Browse files
Files changed (8) hide show
  1. .DS_Store +0 -0
  2. .gitattributes +2 -0
  3. Dockerfile +32 -0
  4. README.md +10 -0
  5. app.py +132 -0
  6. download_models.py +50 -0
  7. requirements.txt +21 -0
  8. utils.py +416 -0
.DS_Store ADDED
Binary file (6.15 kB). View file
 
.gitattributes ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
2
+ *.pt filter=lfs diff=lfs merge=lfs -text
Dockerfile ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9-slim
2
+
3
+ # Install system dependencies
4
+ RUN apt-get update && apt-get install -y \
5
+ libgl1-mesa-glx \
6
+ libglib2.0-0 \
7
+ wget \
8
+ git \
9
+ && rm -rf /var/lib/apt/lists/*
10
+
11
+ WORKDIR /app
12
+
13
+ # Copy requirements and install dependencies
14
+ COPY requirements.txt .
15
+ RUN pip install --no-cache-dir -r requirements.txt
16
+
17
+ # Create directories for models
18
+ RUN mkdir -p weights/icon_detect weights/icon_caption_florence
19
+
20
+ # Copy application code
21
+ COPY main.py utils.py download_models.py ./
22
+
23
+ # Download models during build
24
+ RUN python download_models.py
25
+
26
+ # Set up user for HF Spaces
27
+ RUN useradd -m -u 1000 user
28
+ USER user
29
+ ENV PATH="/home/user/.local/bin:$PATH"
30
+
31
+ # Run the application
32
+ CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
README.md ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Parser Own
3
+ emoji: 💻
4
+ colorFrom: blue
5
+ colorTo: gray
6
+ sdk: docker
7
+ pinned: false
8
+ ---
9
+
10
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, File, UploadFile, HTTPException
2
+ from fastapi.responses import JSONResponse
3
+ from pydantic import BaseModel
4
+ from typing import Optional
5
+ import base64
6
+ import io
7
+ from PIL import Image
8
+ import torch
9
+ import numpy as np
10
+ import os
11
+
12
+ # Existing imports
13
+ import numpy as np
14
+ import torch
15
+ from PIL import Image
16
+ import io
17
+
18
+ from utils import (
19
+ check_ocr_box,
20
+ get_yolo_model,
21
+ get_caption_model_processor,
22
+ get_som_labeled_img,
23
+ )
24
+ import torch
25
+
26
+ # yolo_model = get_yolo_model(model_path='/data/icon_detect/best.pt')
27
+ # caption_model_processor = get_caption_model_processor(model_name="florence2", model_name_or_path="/data/icon_caption_florence")
28
+
29
+ from ultralytics import YOLO
30
+
31
+ # if not os.path.exists("/data/icon_detect"):
32
+ # os.makedirs("/data/icon_detect")
33
+
34
+ try:
35
+ yolo_model = YOLO("weights/icon_detect/best.pt").to("cuda")
36
+ except:
37
+ yolo_model = YOLO("weights/icon_detect/best.pt")
38
+
39
+ from transformers import AutoProcessor, AutoModelForCausalLM
40
+
41
+ processor = AutoProcessor.from_pretrained(
42
+ "microsoft/Florence-2-base", trust_remote_code=True
43
+ )
44
+
45
+ try:
46
+ model = AutoModelForCausalLM.from_pretrained(
47
+ "weights/icon_caption_florence",
48
+ torch_dtype=torch.float16,
49
+ trust_remote_code=True,
50
+ ).to("cuda")
51
+ except:
52
+ model = AutoModelForCausalLM.from_pretrained(
53
+ "weights/icon_caption_florence",
54
+ torch_dtype=torch.float16,
55
+ trust_remote_code=True,
56
+ )
57
+ caption_model_processor = {"processor": processor, "model": model}
58
+ print("finish loading model!!!")
59
+
60
+ app = FastAPI()
61
+
62
+
63
+ class ProcessResponse(BaseModel):
64
+ image: str # Base64 encoded image
65
+ parsed_content_list: str
66
+ label_coordinates: str
67
+
68
+
69
+ def process(
70
+ image_input: Image.Image, box_threshold: float, iou_threshold: float
71
+ ) -> ProcessResponse:
72
+ image_save_path = "imgs/saved_image_demo.png"
73
+ image_input.save(image_save_path)
74
+ image = Image.open(image_save_path)
75
+ box_overlay_ratio = image.size[0] / 3200
76
+ draw_bbox_config = {
77
+ "text_scale": 0.8 * box_overlay_ratio,
78
+ "text_thickness": max(int(2 * box_overlay_ratio), 1),
79
+ "text_padding": max(int(3 * box_overlay_ratio), 1),
80
+ "thickness": max(int(3 * box_overlay_ratio), 1),
81
+ }
82
+
83
+ ocr_bbox_rslt, is_goal_filtered = check_ocr_box(
84
+ image_save_path,
85
+ display_img=False,
86
+ output_bb_format="xyxy",
87
+ goal_filtering=None,
88
+ easyocr_args={"paragraph": False, "text_threshold": 0.9},
89
+ use_paddleocr=True,
90
+ )
91
+ text, ocr_bbox = ocr_bbox_rslt
92
+ dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
93
+ image_save_path,
94
+ yolo_model,
95
+ BOX_TRESHOLD=box_threshold,
96
+ output_coord_in_ratio=True,
97
+ ocr_bbox=ocr_bbox,
98
+ draw_bbox_config=draw_bbox_config,
99
+ caption_model_processor=caption_model_processor,
100
+ ocr_text=text,
101
+ iou_threshold=iou_threshold,
102
+ )
103
+ image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
104
+ print("finish processing")
105
+ parsed_content_list_str = "\n".join(parsed_content_list)
106
+
107
+ # Encode image to base64
108
+ buffered = io.BytesIO()
109
+ image.save(buffered, format="PNG")
110
+ img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
111
+
112
+ return ProcessResponse(
113
+ image=img_str,
114
+ parsed_content_list=str(parsed_content_list_str),
115
+ label_coordinates=str(label_coordinates),
116
+ )
117
+
118
+
119
+ @app.post("/process_image", response_model=ProcessResponse)
120
+ async def process_image(
121
+ image_file: UploadFile = File(...),
122
+ box_threshold: float = 0.05,
123
+ iou_threshold: float = 0.1,
124
+ ):
125
+ try:
126
+ contents = await image_file.read()
127
+ image_input = Image.open(io.BytesIO(contents)).convert("RGB")
128
+ except Exception as e:
129
+ raise HTTPException(status_code=400, detail="Invalid image file")
130
+
131
+ response = process(image_input, box_threshold, iou_threshold)
132
+ return response
download_models.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from huggingface_hub import hf_hub_download
3
+
4
+ def download_models():
5
+ # Create directories if they don't exist
6
+ os.makedirs("weights/icon_detect", exist_ok=True)
7
+ os.makedirs("weights/icon_caption_florence", exist_ok=True)
8
+
9
+ # Files to download from banao-tech/OmniParser
10
+ florence_files = [
11
+ "weights/icon_caption_florence/config.json",
12
+ "weights/icon_caption_florence/generation_config.json",
13
+ "weights/icon_caption_florence/model.safetensors"
14
+ ]
15
+
16
+ yolo_files = [
17
+ "weights/icon_detect/best.pt",
18
+ "weights/icon_detect/model.yaml"
19
+ ]
20
+
21
+ # Download Florence model files
22
+ for file_path in florence_files:
23
+ if not os.path.exists(file_path):
24
+ print(f"Downloading {file_path}...")
25
+ try:
26
+ hf_hub_download(
27
+ repo_id="banao-tech/OmniParser",
28
+ filename=file_path,
29
+ local_dir="."
30
+ )
31
+ print(f"Successfully downloaded {file_path}")
32
+ except Exception as e:
33
+ print(f"Error downloading {file_path}: {str(e)}")
34
+
35
+ # Download YOLO model files
36
+ for file_path in yolo_files:
37
+ if not os.path.exists(file_path):
38
+ print(f"Downloading {file_path}...")
39
+ try:
40
+ hf_hub_download(
41
+ repo_id="banao-tech/OmniParser",
42
+ filename=file_path,
43
+ local_dir="."
44
+ )
45
+ print(f"Successfully downloaded {file_path}")
46
+ except Exception as e:
47
+ print(f"Error downloading {file_path}: {str(e)}")
48
+
49
+ if __name__ == "__main__":
50
+ download_models()
requirements.txt ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ easyocr
3
+ torchvision
4
+ supervision==0.18.0
5
+ openai==1.3.5
6
+ transformers
7
+ ultralytics==8.1.24
8
+ azure-identity
9
+ numpy
10
+ opencv-python
11
+ opencv-python-headless
12
+ gradio
13
+ dill
14
+ accelerate
15
+ timm
16
+ einops==0.8.0
17
+ paddlepaddle
18
+ paddleocr
19
+ fastapi
20
+ uvicorn
21
+ huggingface_hub
utils.py ADDED
@@ -0,0 +1,416 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # from ultralytics import YOLO
2
+ import os
3
+ import io
4
+ import base64
5
+ import time
6
+ from PIL import Image, ImageDraw, ImageFont
7
+ import json
8
+ import requests
9
+ # utility function
10
+ import os
11
+
12
+ import json
13
+ import sys
14
+ import os
15
+ import cv2
16
+ import numpy as np
17
+ # %matplotlib inline
18
+ from matplotlib import pyplot as plt
19
+ import easyocr
20
+ from paddleocr import PaddleOCR
21
+ reader = easyocr.Reader(['en'])
22
+ paddle_ocr = PaddleOCR(
23
+ lang='en', # other lang also available
24
+ use_angle_cls=False,
25
+ use_gpu=False, # using cuda will conflict with pytorch in the same process
26
+ show_log=False,
27
+ max_batch_size=1024,
28
+ use_dilation=True, # improves accuracy
29
+ det_db_score_mode='slow', # improves accuracy
30
+ rec_batch_num=1024)
31
+ import time
32
+ import base64
33
+
34
+ import os
35
+ import ast
36
+ import torch
37
+ from typing import Tuple, List
38
+ from torchvision.ops import box_convert
39
+ import re
40
+ from torchvision.transforms import ToPILImage
41
+ import supervision as sv
42
+ import torchvision.transforms as T
43
+
44
+
45
+ def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
46
+ if not device:
47
+ device = "cuda" if torch.cuda.is_available() else "cpu"
48
+ if model_name == "blip2":
49
+ from transformers import Blip2Processor, Blip2ForConditionalGeneration
50
+ processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
51
+ if device == 'cpu':
52
+ model = Blip2ForConditionalGeneration.from_pretrained(
53
+ model_name_or_path, device_map=None, torch_dtype=torch.float32
54
+ )
55
+ else:
56
+ model = Blip2ForConditionalGeneration.from_pretrained(
57
+ model_name_or_path, device_map=None, torch_dtype=torch.float16
58
+ ).to(device)
59
+ elif model_name == "florence2":
60
+ from transformers import AutoProcessor, AutoModelForCausalLM
61
+ processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
62
+ if device == 'cpu':
63
+ model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True)
64
+ else:
65
+ model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True).to(device)
66
+ return {'model': model.to(device), 'processor': processor}
67
+
68
+
69
+ def get_yolo_model(model_path):
70
+ from ultralytics import YOLO
71
+ # Load the model.
72
+ model = YOLO(model_path)
73
+ return model
74
+
75
+
76
+ @torch.inference_mode()
77
+ def get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=None):
78
+ to_pil = ToPILImage()
79
+ if ocr_bbox:
80
+ non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
81
+ else:
82
+ non_ocr_boxes = filtered_boxes
83
+ croped_pil_image = []
84
+ for i, coord in enumerate(non_ocr_boxes):
85
+ xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
86
+ ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
87
+ cropped_image = image_source[ymin:ymax, xmin:xmax, :]
88
+ croped_pil_image.append(to_pil(cropped_image))
89
+
90
+ model, processor = caption_model_processor['model'], caption_model_processor['processor']
91
+ if not prompt:
92
+ if 'florence' in model.config.name_or_path:
93
+ prompt = "<CAPTION>"
94
+ else:
95
+ prompt = "The image shows"
96
+
97
+ batch_size = 10 # Number of samples per batch
98
+ generated_texts = []
99
+ device = model.device
100
+
101
+ for i in range(0, len(croped_pil_image), batch_size):
102
+ batch = croped_pil_image[i:i+batch_size]
103
+ if model.device.type == 'cuda':
104
+ inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device, dtype=torch.float16)
105
+ else:
106
+ inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
107
+ if 'florence' in model.config.name_or_path:
108
+ generated_ids = model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=1024,num_beams=3, do_sample=False)
109
+ else:
110
+ generated_ids = model.generate(**inputs, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True, num_return_sequences=1) # temperature=0.01, do_sample=True,
111
+ generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
112
+ generated_text = [gen.strip() for gen in generated_text]
113
+ generated_texts.extend(generated_text)
114
+
115
+ return generated_texts
116
+
117
+
118
+
119
+ def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor):
120
+ to_pil = ToPILImage()
121
+ if ocr_bbox:
122
+ non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
123
+ else:
124
+ non_ocr_boxes = filtered_boxes
125
+ croped_pil_image = []
126
+ for i, coord in enumerate(non_ocr_boxes):
127
+ xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
128
+ ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
129
+ cropped_image = image_source[ymin:ymax, xmin:xmax, :]
130
+ croped_pil_image.append(to_pil(cropped_image))
131
+
132
+ model, processor = caption_model_processor['model'], caption_model_processor['processor']
133
+ device = model.device
134
+ messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}]
135
+ prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
136
+
137
+ batch_size = 5 # Number of samples per batch
138
+ generated_texts = []
139
+
140
+ for i in range(0, len(croped_pil_image), batch_size):
141
+ images = croped_pil_image[i:i+batch_size]
142
+ image_inputs = [processor.image_processor(x, return_tensors="pt") for x in images]
143
+ inputs ={'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
144
+ texts = [prompt] * len(images)
145
+ for i, txt in enumerate(texts):
146
+ input = processor._convert_images_texts_to_inputs(image_inputs[i], txt, return_tensors="pt")
147
+ inputs['input_ids'].append(input['input_ids'])
148
+ inputs['attention_mask'].append(input['attention_mask'])
149
+ inputs['pixel_values'].append(input['pixel_values'])
150
+ inputs['image_sizes'].append(input['image_sizes'])
151
+ max_len = max([x.shape[1] for x in inputs['input_ids']])
152
+ for i, v in enumerate(inputs['input_ids']):
153
+ inputs['input_ids'][i] = torch.cat([processor.tokenizer.pad_token_id * torch.ones(1, max_len - v.shape[1], dtype=torch.long), v], dim=1)
154
+ inputs['attention_mask'][i] = torch.cat([torch.zeros(1, max_len - v.shape[1], dtype=torch.long), inputs['attention_mask'][i]], dim=1)
155
+ inputs_cat = {k: torch.concatenate(v).to(device) for k, v in inputs.items()}
156
+
157
+ generation_args = {
158
+ "max_new_tokens": 25,
159
+ "temperature": 0.01,
160
+ "do_sample": False,
161
+ }
162
+ generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
163
+ # # remove input tokens
164
+ generate_ids = generate_ids[:, inputs_cat['input_ids'].shape[1]:]
165
+ response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
166
+ response = [res.strip('\n').strip() for res in response]
167
+ generated_texts.extend(response)
168
+
169
+ return generated_texts
170
+
171
+ def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
172
+ assert ocr_bbox is None or isinstance(ocr_bbox, List)
173
+
174
+ def box_area(box):
175
+ return (box[2] - box[0]) * (box[3] - box[1])
176
+
177
+ def intersection_area(box1, box2):
178
+ x1 = max(box1[0], box2[0])
179
+ y1 = max(box1[1], box2[1])
180
+ x2 = min(box1[2], box2[2])
181
+ y2 = min(box1[3], box2[3])
182
+ return max(0, x2 - x1) * max(0, y2 - y1)
183
+
184
+ def IoU(box1, box2):
185
+ intersection = intersection_area(box1, box2)
186
+ union = box_area(box1) + box_area(box2) - intersection + 1e-6
187
+ if box_area(box1) > 0 and box_area(box2) > 0:
188
+ ratio1 = intersection / box_area(box1)
189
+ ratio2 = intersection / box_area(box2)
190
+ else:
191
+ ratio1, ratio2 = 0, 0
192
+ return max(intersection / union, ratio1, ratio2)
193
+
194
+ boxes = boxes.tolist()
195
+ filtered_boxes = []
196
+ if ocr_bbox:
197
+ filtered_boxes.extend(ocr_bbox)
198
+ # print('ocr_bbox!!!', ocr_bbox)
199
+ for i, box1 in enumerate(boxes):
200
+ # if not any(IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2) for j, box2 in enumerate(boxes) if i != j):
201
+ is_valid_box = True
202
+ for j, box2 in enumerate(boxes):
203
+ if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
204
+ is_valid_box = False
205
+ break
206
+ if is_valid_box:
207
+ # add the following 2 lines to include ocr bbox
208
+ if ocr_bbox:
209
+ if not any(IoU(box1, box3) > iou_threshold for k, box3 in enumerate(ocr_bbox)):
210
+ filtered_boxes.append(box1)
211
+ else:
212
+ filtered_boxes.append(box1)
213
+ return torch.tensor(filtered_boxes)
214
+
215
+ def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
216
+ transform = T.Compose(
217
+ [
218
+ T.RandomResize([800], max_size=1333),
219
+ T.ToTensor(),
220
+ T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
221
+ ]
222
+ )
223
+ image_source = Image.open(image_path).convert("RGB")
224
+ image = np.asarray(image_source)
225
+ image_transformed, _ = transform(image_source, None)
226
+ return image, image_transformed
227
+
228
+
229
+ def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], text_scale: float,
230
+ text_padding=5, text_thickness=2, thickness=3) -> np.ndarray:
231
+ """
232
+ This function annotates an image with bounding boxes and labels.
233
+
234
+ Parameters:
235
+ image_source (np.ndarray): The source image to be annotated.
236
+ boxes (torch.Tensor): A tensor containing bounding box coordinates. in cxcywh format, pixel scale
237
+ logits (torch.Tensor): A tensor containing confidence scores for each bounding box.
238
+ phrases (List[str]): A list of labels for each bounding box.
239
+ text_scale (float): The scale of the text to be displayed. 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
240
+
241
+ Returns:
242
+ np.ndarray: The annotated image.
243
+ """
244
+ h, w, _ = image_source.shape
245
+ boxes = boxes * torch.Tensor([w, h, w, h])
246
+ xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
247
+ xywh = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xywh").numpy()
248
+ detections = sv.Detections(xyxy=xyxy)
249
+
250
+ labels = [f"{phrase}" for phrase in range(boxes.shape[0])]
251
+
252
+ from util.box_annotator import BoxAnnotator
253
+ box_annotator = BoxAnnotator(text_scale=text_scale, text_padding=text_padding,text_thickness=text_thickness,thickness=thickness) # 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
254
+ annotated_frame = image_source.copy()
255
+ annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w,h))
256
+
257
+ label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
258
+ return annotated_frame, label_coordinates
259
+
260
+
261
+ def predict(model, image, caption, box_threshold, text_threshold):
262
+ """ Use huggingface model to replace the original model
263
+ """
264
+ model, processor = model['model'], model['processor']
265
+ device = model.device
266
+
267
+ inputs = processor(images=image, text=caption, return_tensors="pt").to(device)
268
+ with torch.no_grad():
269
+ outputs = model(**inputs)
270
+
271
+ results = processor.post_process_grounded_object_detection(
272
+ outputs,
273
+ inputs.input_ids,
274
+ box_threshold=box_threshold, # 0.4,
275
+ text_threshold=text_threshold, # 0.3,
276
+ target_sizes=[image.size[::-1]]
277
+ )[0]
278
+ boxes, logits, phrases = results["boxes"], results["scores"], results["labels"]
279
+ return boxes, logits, phrases
280
+
281
+
282
+ def predict_yolo(model, image_path, box_threshold):
283
+ """ Use huggingface model to replace the original model
284
+ """
285
+ # model = model['model']
286
+
287
+ result = model.predict(
288
+ source=image_path,
289
+ conf=box_threshold,
290
+ # iou=0.5, # default 0.7
291
+ )
292
+ boxes = result[0].boxes.xyxy#.tolist() # in pixel space
293
+ conf = result[0].boxes.conf
294
+ phrases = [str(i) for i in range(len(boxes))]
295
+
296
+ return boxes, conf, phrases
297
+
298
+
299
+ def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None):
300
+ """ ocr_bbox: list of xyxy format bbox
301
+ """
302
+ TEXT_PROMPT = "clickable buttons on the screen"
303
+ # BOX_TRESHOLD = 0.02 # 0.05/0.02 for web and 0.1 for mobile
304
+ TEXT_TRESHOLD = 0.01 # 0.9 # 0.01
305
+ image_source = Image.open(img_path).convert("RGB")
306
+ w, h = image_source.size
307
+ # import pdb; pdb.set_trace()
308
+ if False: # TODO
309
+ xyxy, logits, phrases = predict(model=model, image=image_source, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD)
310
+ else:
311
+ xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD)
312
+ xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device)
313
+ image_source = np.asarray(image_source)
314
+ phrases = [str(i) for i in range(len(phrases))]
315
+
316
+ # annotate the image with labels
317
+ h, w, _ = image_source.shape
318
+ if ocr_bbox:
319
+ ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
320
+ ocr_bbox=ocr_bbox.tolist()
321
+ else:
322
+ print('no ocr bbox!!!')
323
+ ocr_bbox = None
324
+ filtered_boxes = remove_overlap(boxes=xyxy, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox)
325
+
326
+ # get parsed icon local semantics
327
+ if use_local_semantics:
328
+ caption_model = caption_model_processor['model']
329
+ if 'phi3_v' in caption_model.config.model_type:
330
+ parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor)
331
+ else:
332
+ parsed_content_icon = get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=prompt)
333
+ ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
334
+ icon_start = len(ocr_text)
335
+ parsed_content_icon_ls = []
336
+ for i, txt in enumerate(parsed_content_icon):
337
+ parsed_content_icon_ls.append(f"Icon Box ID {str(i+icon_start)}: {txt}")
338
+ parsed_content_merged = ocr_text + parsed_content_icon_ls
339
+ else:
340
+ ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
341
+ parsed_content_merged = ocr_text
342
+
343
+ filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
344
+
345
+ phrases = [i for i in range(len(filtered_boxes))]
346
+
347
+ # draw boxes
348
+ if draw_bbox_config:
349
+ annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
350
+ else:
351
+ annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding)
352
+
353
+ pil_img = Image.fromarray(annotated_frame)
354
+ buffered = io.BytesIO()
355
+ pil_img.save(buffered, format="PNG")
356
+ encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
357
+ if output_coord_in_ratio:
358
+ # h, w, _ = image_source.shape
359
+ label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()}
360
+ assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
361
+
362
+ return encoded_image, label_coordinates, parsed_content_merged
363
+
364
+
365
+ def get_xywh(input):
366
+ x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1]
367
+ x, y, w, h = int(x), int(y), int(w), int(h)
368
+ return x, y, w, h
369
+
370
+ def get_xyxy(input):
371
+ x, y, xp, yp = input[0][0], input[0][1], input[2][0], input[2][1]
372
+ x, y, xp, yp = int(x), int(y), int(xp), int(yp)
373
+ return x, y, xp, yp
374
+
375
+ def get_xywh_yolo(input):
376
+ x, y, w, h = input[0], input[1], input[2] - input[0], input[3] - input[1]
377
+ x, y, w, h = int(x), int(y), int(w), int(h)
378
+ return x, y, w, h
379
+
380
+
381
+
382
+ def check_ocr_box(image_path, display_img = True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None, use_paddleocr=False):
383
+ if use_paddleocr:
384
+ result = paddle_ocr.ocr(image_path, cls=False)[0]
385
+ coord = [item[0] for item in result]
386
+ text = [item[1][0] for item in result]
387
+ else: # EasyOCR
388
+ if easyocr_args is None:
389
+ easyocr_args = {}
390
+ result = reader.readtext(image_path, **easyocr_args)
391
+ # print('goal filtering pred:', result[-5:])
392
+ coord = [item[0] for item in result]
393
+ text = [item[1] for item in result]
394
+ # read the image using cv2
395
+ if display_img:
396
+ opencv_img = cv2.imread(image_path)
397
+ opencv_img = cv2.cvtColor(opencv_img, cv2.COLOR_RGB2BGR)
398
+ bb = []
399
+ for item in coord:
400
+ x, y, a, b = get_xywh(item)
401
+ # print(x, y, a, b)
402
+ bb.append((x, y, a, b))
403
+ cv2.rectangle(opencv_img, (x, y), (x+a, y+b), (0, 255, 0), 2)
404
+
405
+ # Display the image
406
+ plt.imshow(opencv_img)
407
+ else:
408
+ if output_bb_format == 'xywh':
409
+ bb = [get_xywh(item) for item in coord]
410
+ elif output_bb_format == 'xyxy':
411
+ bb = [get_xyxy(item) for item in coord]
412
+ # print('bounding box!!!', bb)
413
+ return (text, bb), goal_filtering
414
+
415
+
416
+