gaoyu1314 commited on
Commit
8d999d8
·
verified ·
1 Parent(s): 000d129

Delete utils.py

Browse files
Files changed (1) hide show
  1. utils.py +0 -540
utils.py DELETED
@@ -1,540 +0,0 @@
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
- from openai import AzureOpenAI
12
-
13
- import json
14
- import sys
15
- import os
16
- import cv2
17
- import numpy as np
18
- # %matplotlib inline
19
- from matplotlib import pyplot as plt
20
- import easyocr
21
- from paddleocr import PaddleOCR
22
- reader = easyocr.Reader(['en'])
23
- paddle_ocr = PaddleOCR(
24
- lang='en', # other lang also available
25
- use_angle_cls=False,
26
- use_gpu=False, # using cuda will conflict with pytorch in the same process
27
- show_log=False,
28
- max_batch_size=1024,
29
- use_dilation=True, # improves accuracy
30
- det_db_score_mode='slow', # improves accuracy
31
- rec_batch_num=1024)
32
- import time
33
- import base64
34
-
35
- import os
36
- import ast
37
- import torch
38
- from typing import Tuple, List, Union
39
- from torchvision.ops import box_convert
40
- import re
41
- from torchvision.transforms import ToPILImage
42
- import supervision as sv
43
- import torchvision.transforms as T
44
- from util.box_annotator import BoxAnnotator
45
-
46
-
47
- def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
48
- if not device:
49
- device = "cuda" if torch.cuda.is_available() else "cpu"
50
- if model_name == "blip2":
51
- from transformers import Blip2Processor, Blip2ForConditionalGeneration
52
- processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
53
- if device == 'cpu':
54
- model = Blip2ForConditionalGeneration.from_pretrained(
55
- model_name_or_path, device_map=None, torch_dtype=torch.float32
56
- )
57
- else:
58
- model = Blip2ForConditionalGeneration.from_pretrained(
59
- model_name_or_path, device_map=None, torch_dtype=torch.float16
60
- ).to(device)
61
- elif model_name == "florence2":
62
- from transformers import AutoProcessor, AutoModelForCausalLM
63
- processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
64
- if device == 'cpu':
65
- model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True)
66
- else:
67
- model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True).to(device)
68
- return {'model': model.to(device), 'processor': processor}
69
-
70
-
71
- def get_yolo_model(model_path):
72
- from ultralytics import YOLO
73
- # Load the model.
74
- model = YOLO(model_path)
75
- return model
76
-
77
-
78
- @torch.inference_mode()
79
- def get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=None, batch_size=128):
80
- # Number of samples per batch, --> 128 roughly takes 4 GB of GPU memory for florence v2 model
81
- to_pil = ToPILImage()
82
- if starting_idx:
83
- non_ocr_boxes = filtered_boxes[starting_idx:]
84
- else:
85
- non_ocr_boxes = filtered_boxes
86
- croped_pil_image = []
87
- for i, coord in enumerate(non_ocr_boxes):
88
- try:
89
- xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
90
- ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
91
- cropped_image = image_source[ymin:ymax, xmin:xmax, :]
92
- cropped_image = cv2.resize(cropped_image, (64, 64))
93
- croped_pil_image.append(to_pil(cropped_image))
94
- except:
95
- continue
96
-
97
- model, processor = caption_model_processor['model'], caption_model_processor['processor']
98
- if not prompt:
99
- if 'florence' in model.config.name_or_path:
100
- prompt = "<CAPTION>"
101
- else:
102
- prompt = "The image shows"
103
-
104
- generated_texts = []
105
- device = model.device
106
- for i in range(0, len(croped_pil_image), batch_size):
107
- start = time.time()
108
- batch = croped_pil_image[i:i+batch_size]
109
- t1 = time.time()
110
- if model.device.type == 'cuda':
111
- inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt", do_resize=False).to(device=device, dtype=torch.float16)
112
- else:
113
- inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
114
- if 'florence' in model.config.name_or_path:
115
- generated_ids = model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=20,num_beams=1, do_sample=False)
116
- else:
117
- 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,
118
- generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
119
- generated_text = [gen.strip() for gen in generated_text]
120
- generated_texts.extend(generated_text)
121
-
122
- return generated_texts
123
-
124
-
125
-
126
- def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor):
127
- to_pil = ToPILImage()
128
- if ocr_bbox:
129
- non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
130
- else:
131
- non_ocr_boxes = filtered_boxes
132
- croped_pil_image = []
133
- for i, coord in enumerate(non_ocr_boxes):
134
- xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
135
- ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
136
- cropped_image = image_source[ymin:ymax, xmin:xmax, :]
137
- croped_pil_image.append(to_pil(cropped_image))
138
-
139
- model, processor = caption_model_processor['model'], caption_model_processor['processor']
140
- device = model.device
141
- messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}]
142
- prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
143
-
144
- batch_size = 5 # Number of samples per batch
145
- generated_texts = []
146
-
147
- for i in range(0, len(croped_pil_image), batch_size):
148
- images = croped_pil_image[i:i+batch_size]
149
- image_inputs = [processor.image_processor(x, return_tensors="pt") for x in images]
150
- inputs ={'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
151
- texts = [prompt] * len(images)
152
- for i, txt in enumerate(texts):
153
- input = processor._convert_images_texts_to_inputs(image_inputs[i], txt, return_tensors="pt")
154
- inputs['input_ids'].append(input['input_ids'])
155
- inputs['attention_mask'].append(input['attention_mask'])
156
- inputs['pixel_values'].append(input['pixel_values'])
157
- inputs['image_sizes'].append(input['image_sizes'])
158
- max_len = max([x.shape[1] for x in inputs['input_ids']])
159
- for i, v in enumerate(inputs['input_ids']):
160
- 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)
161
- inputs['attention_mask'][i] = torch.cat([torch.zeros(1, max_len - v.shape[1], dtype=torch.long), inputs['attention_mask'][i]], dim=1)
162
- inputs_cat = {k: torch.concatenate(v).to(device) for k, v in inputs.items()}
163
-
164
- generation_args = {
165
- "max_new_tokens": 25,
166
- "temperature": 0.01,
167
- "do_sample": False,
168
- }
169
- generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
170
- # # remove input tokens
171
- generate_ids = generate_ids[:, inputs_cat['input_ids'].shape[1]:]
172
- response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
173
- response = [res.strip('\n').strip() for res in response]
174
- generated_texts.extend(response)
175
-
176
- return generated_texts
177
-
178
- def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
179
- assert ocr_bbox is None or isinstance(ocr_bbox, List)
180
-
181
- def box_area(box):
182
- return (box[2] - box[0]) * (box[3] - box[1])
183
-
184
- def intersection_area(box1, box2):
185
- x1 = max(box1[0], box2[0])
186
- y1 = max(box1[1], box2[1])
187
- x2 = min(box1[2], box2[2])
188
- y2 = min(box1[3], box2[3])
189
- return max(0, x2 - x1) * max(0, y2 - y1)
190
-
191
- def IoU(box1, box2):
192
- intersection = intersection_area(box1, box2)
193
- union = box_area(box1) + box_area(box2) - intersection + 1e-6
194
- if box_area(box1) > 0 and box_area(box2) > 0:
195
- ratio1 = intersection / box_area(box1)
196
- ratio2 = intersection / box_area(box2)
197
- else:
198
- ratio1, ratio2 = 0, 0
199
- return max(intersection / union, ratio1, ratio2)
200
-
201
- def is_inside(box1, box2):
202
- # return box1[0] >= box2[0] and box1[1] >= box2[1] and box1[2] <= box2[2] and box1[3] <= box2[3]
203
- intersection = intersection_area(box1, box2)
204
- ratio1 = intersection / box_area(box1)
205
- return ratio1 > 0.95
206
-
207
- boxes = boxes.tolist()
208
- filtered_boxes = []
209
- if ocr_bbox:
210
- filtered_boxes.extend(ocr_bbox)
211
- # print('ocr_bbox!!!', ocr_bbox)
212
- for i, box1 in enumerate(boxes):
213
- # if not any(IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2) for j, box2 in enumerate(boxes) if i != j):
214
- is_valid_box = True
215
- for j, box2 in enumerate(boxes):
216
- # keep the smaller box
217
- if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
218
- is_valid_box = False
219
- break
220
- if is_valid_box:
221
- # add the following 2 lines to include ocr bbox
222
- if ocr_bbox:
223
- # only add the box if it does not overlap with any ocr bbox
224
- if not any(IoU(box1, box3) > iou_threshold and not is_inside(box1, box3) for k, box3 in enumerate(ocr_bbox)):
225
- filtered_boxes.append(box1)
226
- else:
227
- filtered_boxes.append(box1)
228
- return torch.tensor(filtered_boxes)
229
-
230
-
231
- def remove_overlap_new(boxes, iou_threshold, ocr_bbox=None):
232
- '''
233
- ocr_bbox format: [{'type': 'text', 'bbox':[x,y], 'interactivity':False, 'content':str }, ...]
234
- boxes format: [{'type': 'icon', 'bbox':[x,y], 'interactivity':True, 'content':None }, ...]
235
-
236
- '''
237
- assert ocr_bbox is None or isinstance(ocr_bbox, List)
238
-
239
- def box_area(box):
240
- return (box[2] - box[0]) * (box[3] - box[1])
241
-
242
- def intersection_area(box1, box2):
243
- x1 = max(box1[0], box2[0])
244
- y1 = max(box1[1], box2[1])
245
- x2 = min(box1[2], box2[2])
246
- y2 = min(box1[3], box2[3])
247
- return max(0, x2 - x1) * max(0, y2 - y1)
248
-
249
- def IoU(box1, box2):
250
- intersection = intersection_area(box1, box2)
251
- union = box_area(box1) + box_area(box2) - intersection + 1e-6
252
- if box_area(box1) > 0 and box_area(box2) > 0:
253
- ratio1 = intersection / box_area(box1)
254
- ratio2 = intersection / box_area(box2)
255
- else:
256
- ratio1, ratio2 = 0, 0
257
- return max(intersection / union, ratio1, ratio2)
258
-
259
- def is_inside(box1, box2):
260
- # return box1[0] >= box2[0] and box1[1] >= box2[1] and box1[2] <= box2[2] and box1[3] <= box2[3]
261
- intersection = intersection_area(box1, box2)
262
- ratio1 = intersection / box_area(box1)
263
- return ratio1 > 0.80
264
-
265
- # boxes = boxes.tolist()
266
- filtered_boxes = []
267
- if ocr_bbox:
268
- filtered_boxes.extend(ocr_bbox)
269
- # print('ocr_bbox!!!', ocr_bbox)
270
- for i, box1_elem in enumerate(boxes):
271
- box1 = box1_elem['bbox']
272
- is_valid_box = True
273
- for j, box2_elem in enumerate(boxes):
274
- # keep the smaller box
275
- box2 = box2_elem['bbox']
276
- if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
277
- is_valid_box = False
278
- break
279
- if is_valid_box:
280
- if ocr_bbox:
281
- # keep yolo boxes + prioritize ocr label
282
- box_added = False
283
- ocr_labels = ''
284
- for box3_elem in ocr_bbox:
285
- if not box_added:
286
- box3 = box3_elem['bbox']
287
- if is_inside(box3, box1): # ocr inside icon
288
- # box_added = True
289
- # delete the box3_elem from ocr_bbox
290
- try:
291
- # gather all ocr labels
292
- ocr_labels += box3_elem['content'] + ' '
293
- filtered_boxes.remove(box3_elem)
294
- except:
295
- continue
296
- # break
297
- elif is_inside(box1, box3): # icon inside ocr, don't added this icon box, no need to check other ocr bbox bc no overlap between ocr bbox, icon can only be in one ocr box
298
- box_added = True
299
- break
300
- else:
301
- continue
302
- if not box_added:
303
- if ocr_labels:
304
- filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': ocr_labels, 'source':'box_yolo_content_ocr'})
305
- else:
306
- filtered_boxes.append({'type': 'icon', 'bbox': box1_elem['bbox'], 'interactivity': True, 'content': None, 'source':'box_yolo_content_yolo'})
307
- else:
308
- filtered_boxes.append(box1)
309
- return filtered_boxes # torch.tensor(filtered_boxes)
310
-
311
-
312
- def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
313
- transform = T.Compose(
314
- [
315
- T.RandomResize([800], max_size=1333),
316
- T.ToTensor(),
317
- T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
318
- ]
319
- )
320
- image_source = Image.open(image_path).convert("RGB")
321
- image = np.asarray(image_source)
322
- image_transformed, _ = transform(image_source, None)
323
- return image, image_transformed
324
-
325
-
326
- def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], text_scale: float,
327
- text_padding=5, text_thickness=2, thickness=3) -> np.ndarray:
328
- """
329
- This function annotates an image with bounding boxes and labels.
330
-
331
- Parameters:
332
- image_source (np.ndarray): The source image to be annotated.
333
- boxes (torch.Tensor): A tensor containing bounding box coordinates. in cxcywh format, pixel scale
334
- logits (torch.Tensor): A tensor containing confidence scores for each bounding box.
335
- phrases (List[str]): A list of labels for each bounding box.
336
- text_scale (float): The scale of the text to be displayed. 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
337
-
338
- Returns:
339
- np.ndarray: The annotated image.
340
- """
341
- h, w, _ = image_source.shape
342
- boxes = boxes * torch.Tensor([w, h, w, h])
343
- xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
344
- xywh = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xywh").numpy()
345
- detections = sv.Detections(xyxy=xyxy)
346
-
347
- labels = [f"{phrase}" for phrase in range(boxes.shape[0])]
348
-
349
- 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
350
- annotated_frame = image_source.copy()
351
- annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w,h))
352
-
353
- label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
354
- return annotated_frame, label_coordinates
355
-
356
-
357
- def predict(model, image, caption, box_threshold, text_threshold):
358
- """ Use huggingface model to replace the original model
359
- """
360
- model, processor = model['model'], model['processor']
361
- device = model.device
362
-
363
- inputs = processor(images=image, text=caption, return_tensors="pt").to(device)
364
- with torch.no_grad():
365
- outputs = model(**inputs)
366
-
367
- results = processor.post_process_grounded_object_detection(
368
- outputs,
369
- inputs.input_ids,
370
- box_threshold=box_threshold, # 0.4,
371
- text_threshold=text_threshold, # 0.3,
372
- target_sizes=[image.size[::-1]]
373
- )[0]
374
- boxes, logits, phrases = results["boxes"], results["scores"], results["labels"]
375
- return boxes, logits, phrases
376
-
377
-
378
- def predict_yolo(model, image, box_threshold, imgsz, scale_img, iou_threshold=0.7):
379
- """ Use huggingface model to replace the original model
380
- """
381
- # model = model['model']
382
- if scale_img:
383
- result = model.predict(
384
- source=image,
385
- conf=box_threshold,
386
- imgsz=imgsz,
387
- iou=iou_threshold, # default 0.7
388
- )
389
- else:
390
- result = model.predict(
391
- source=image,
392
- conf=box_threshold,
393
- iou=iou_threshold, # default 0.7
394
- )
395
- boxes = result[0].boxes.xyxy#.tolist() # in pixel space
396
- conf = result[0].boxes.conf
397
- phrases = [str(i) for i in range(len(boxes))]
398
-
399
- return boxes, conf, phrases
400
-
401
- def int_box_area(box, w, h):
402
- x1, y1, x2, y2 = box
403
- int_box = [int(x1*w), int(y1*h), int(x2*w), int(y2*h)]
404
- area = (int_box[2] - int_box[0]) * (int_box[3] - int_box[1])
405
- return area
406
-
407
- def get_som_labeled_img(image_source: Union[str, Image.Image], 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, scale_img=False, imgsz=None, batch_size=128):
408
- """Process either an image path or Image object
409
-
410
- Args:
411
- image_source: Either a file path (str) or PIL Image object
412
- ...
413
- """
414
- if isinstance(image_source, str):
415
- image_source = Image.open(image_source)
416
- image_source = image_source.convert("RGB") # for CLIP
417
- w, h = image_source.size
418
- if not imgsz:
419
- imgsz = (h, w)
420
- # print('image size:', w, h)
421
- xyxy, logits, phrases = predict_yolo(model=model, image=image_source, box_threshold=BOX_TRESHOLD, imgsz=imgsz, scale_img=scale_img, iou_threshold=0.1)
422
- xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device)
423
- image_source = np.asarray(image_source)
424
- phrases = [str(i) for i in range(len(phrases))]
425
-
426
- # annotate the image with labels
427
- if ocr_bbox:
428
- ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
429
- ocr_bbox=ocr_bbox.tolist()
430
- else:
431
- print('no ocr bbox!!!')
432
- ocr_bbox = None
433
-
434
- ocr_bbox_elem = [{'type': 'text', 'bbox':box, 'interactivity':False, 'content':txt, 'source': 'box_ocr_content_ocr'} for box, txt in zip(ocr_bbox, ocr_text) if int_box_area(box, w, h) > 0]
435
- xyxy_elem = [{'type': 'icon', 'bbox':box, 'interactivity':True, 'content':None} for box in xyxy.tolist() if int_box_area(box, w, h) > 0]
436
- filtered_boxes = remove_overlap_new(boxes=xyxy_elem, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox_elem)
437
-
438
- # sort the filtered_boxes so that the one with 'content': None is at the end, and get the index of the first 'content': None
439
- filtered_boxes_elem = sorted(filtered_boxes, key=lambda x: x['content'] is None)
440
- # get the index of the first 'content': None
441
- starting_idx = next((i for i, box in enumerate(filtered_boxes_elem) if box['content'] is None), -1)
442
- filtered_boxes = torch.tensor([box['bbox'] for box in filtered_boxes_elem])
443
- print('len(filtered_boxes):', len(filtered_boxes), starting_idx)
444
-
445
- # get parsed icon local semantics
446
- time1 = time.time()
447
- if use_local_semantics:
448
- caption_model = caption_model_processor['model']
449
- if 'phi3_v' in caption_model.config.model_type:
450
- parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor)
451
- else:
452
- parsed_content_icon = get_parsed_content_icon(filtered_boxes, starting_idx, image_source, caption_model_processor, prompt=prompt,batch_size=batch_size)
453
- ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
454
- icon_start = len(ocr_text)
455
- parsed_content_icon_ls = []
456
- # fill the filtered_boxes_elem None content with parsed_content_icon in order
457
- for i, box in enumerate(filtered_boxes_elem):
458
- if box['content'] is None:
459
- box['content'] = parsed_content_icon.pop(0)
460
- for i, txt in enumerate(parsed_content_icon):
461
- parsed_content_icon_ls.append(f"Icon Box ID {str(i+icon_start)}: {txt}")
462
- parsed_content_merged = ocr_text + parsed_content_icon_ls
463
- else:
464
- ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
465
- parsed_content_merged = ocr_text
466
- print('time to get parsed content:', time.time()-time1)
467
-
468
- filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
469
-
470
- phrases = [i for i in range(len(filtered_boxes))]
471
-
472
- # draw boxes
473
- if draw_bbox_config:
474
- annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
475
- else:
476
- annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding)
477
-
478
- pil_img = Image.fromarray(annotated_frame)
479
- buffered = io.BytesIO()
480
- pil_img.save(buffered, format="PNG")
481
- encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
482
- if output_coord_in_ratio:
483
- label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()}
484
- assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
485
-
486
- return encoded_image, label_coordinates, filtered_boxes_elem
487
-
488
-
489
- def get_xywh(input):
490
- x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1]
491
- x, y, w, h = int(x), int(y), int(w), int(h)
492
- return x, y, w, h
493
-
494
- def get_xyxy(input):
495
- x, y, xp, yp = input[0][0], input[0][1], input[2][0], input[2][1]
496
- x, y, xp, yp = int(x), int(y), int(xp), int(yp)
497
- return x, y, xp, yp
498
-
499
- def get_xywh_yolo(input):
500
- x, y, w, h = input[0], input[1], input[2] - input[0], input[3] - input[1]
501
- x, y, w, h = int(x), int(y), int(w), int(h)
502
- return x, y, w, h
503
-
504
- def check_ocr_box(image_source: Union[str, Image.Image], display_img = True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None, use_paddleocr=False):
505
- if isinstance(image_source, str):
506
- image_source = Image.open(image_source)
507
- if image_source.mode == 'RGBA':
508
- # Convert RGBA to RGB to avoid alpha channel issues
509
- image_source = image_source.convert('RGB')
510
- image_np = np.array(image_source)
511
- w, h = image_source.size
512
- if use_paddleocr:
513
- if easyocr_args is None:
514
- text_threshold = 0.5
515
- else:
516
- text_threshold = easyocr_args['text_threshold']
517
- result = paddle_ocr.ocr(image_np, cls=False)[0]
518
- coord = [item[0] for item in result if item[1][1] > text_threshold]
519
- text = [item[1][0] for item in result if item[1][1] > text_threshold]
520
- else: # EasyOCR
521
- if easyocr_args is None:
522
- easyocr_args = {}
523
- result = reader.readtext(image_np, **easyocr_args)
524
- coord = [item[0] for item in result]
525
- text = [item[1] for item in result]
526
- if display_img:
527
- opencv_img = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
528
- bb = []
529
- for item in coord:
530
- x, y, a, b = get_xywh(item)
531
- bb.append((x, y, a, b))
532
- cv2.rectangle(opencv_img, (x, y), (x+a, y+b), (0, 255, 0), 2)
533
- # matplotlib expects RGB
534
- plt.imshow(cv2.cvtColor(opencv_img, cv2.COLOR_BGR2RGB))
535
- else:
536
- if output_bb_format == 'xywh':
537
- bb = [get_xywh(item) for item in coord]
538
- elif output_bb_format == 'xyxy':
539
- bb = [get_xyxy(item) for item in coord]
540
- return (text, bb), goal_filtering