diff --git a/requirements.txt b/requirements.txt index ef8bcc46..aa758eeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,19 +3,33 @@ easyocr torchvision supervision==0.18.0 openai==1.3.5 -transformers + +# Fix for Florence2 Crash +transformers==4.40.0 huggingface_hub + ultralytics==8.3.70 azure-identity numpy==1.26.4 -opencv-python + +# If you need an interactive, desktop-level GUI features, use 'opencv-python' and comment out opencv-python-headless, both of them contain the same core file. + +# Includes GUI dependencies (like Qt) +#opencv-python + +# Removes the GUI dependencies. This is smaller and designed for server environments. opencv-python-headless + gradio dill accelerate timm einops==0.8.0 + +# To use the model on a GPU, change paddlepaddle to: paddlepaddle-gpu paddlepaddle +# paddleocr won't be able to find the GPU since paddlepaddle is a CPU-only library. + paddleocr ruff==0.6.7 pre-commit==3.8.0 @@ -30,4 +44,6 @@ google-auth<3,>=2 screeninfo uiautomation dashscope -groq \ No newline at end of file + +groq + diff --git a/util/utils.py b/util/utils.py index 22c05084..38f14ac0 100644 --- a/util/utils.py +++ b/util/utils.py @@ -22,13 +22,7 @@ reader = easyocr.Reader(['en']) paddle_ocr = PaddleOCR( lang='en', # other lang also available - use_angle_cls=False, - use_gpu=False, # using cuda will conflict with pytorch in the same process - show_log=False, - max_batch_size=1024, - use_dilation=True, # improves accuracy - det_db_score_mode='slow', # improves accuracy - rec_batch_num=1024) +) import time import base64 @@ -42,7 +36,7 @@ from torchvision.transforms import ToPILImage import supervision as sv import torchvision.transforms as T -from util.box_annotator import BoxAnnotator +from .box_annotator import BoxAnnotator def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None): @@ -440,15 +434,29 @@ def get_som_labeled_img(image_source: Union[str, Image.Image], model=None, BOX_T else: print('no ocr bbox!!!') ocr_bbox = None - + if ocr_bbox is None: ocr_bbox = [] + if ocr_text is None: ocr_text = [] 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] xyxy_elem = [{'type': 'icon', 'bbox':box, 'interactivity':True, 'content':None} for box in xyxy.tolist() if int_box_area(box, w, h) > 0] filtered_boxes = remove_overlap_new(boxes=xyxy_elem, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox_elem) # sort the filtered_boxes so that the one with 'content': None is at the end, and get the index of the first 'content': None - filtered_boxes_elem = sorted(filtered_boxes, key=lambda x: x['content'] is None) + # filtered_boxes_elem = sorted(filtered_boxes, key=lambda x: x['content'] is None) + safe_boxes = [] + for item in filtered_boxes: + # If the item is just a list of coordinates (the bug), wrap it in a dict + if isinstance(item, (list, tuple)): + safe_boxes.append({'type': 'icon', 'bbox': list(item), 'interactivity': True, 'content': None, 'source': 'box_yolo_content_yolo'}) + else: + # Ensure required keys exist so downstream code can safely index + if isinstance(item, dict): + item.setdefault('content', None) + safe_boxes.append(item) + + # Now sort the safe list + filtered_boxes_elem = sorted(safe_boxes, key=lambda x: x.get('content') is None) # get the index of the first 'content': None - starting_idx = next((i for i, box in enumerate(filtered_boxes_elem) if box['content'] is None), -1) + starting_idx = next((i for i, box in enumerate(filtered_boxes_elem) if box.get('content') is None), len(filtered_boxes_elem)) filtered_boxes = torch.tensor([box['bbox'] for box in filtered_boxes_elem]) print('len(filtered_boxes):', len(filtered_boxes), starting_idx) @@ -484,18 +492,18 @@ def get_som_labeled_img(image_source: Union[str, Image.Image], model=None, BOX_T annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config) else: annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding) - + pil_img = Image.fromarray(annotated_frame) + buffered = io.BytesIO() pil_img.save(buffered, format="PNG") encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii') + if output_coord_in_ratio: label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()} assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0] - + return encoded_image, label_coordinates, filtered_boxes_elem - - def get_xywh(input): x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1] x, y, w, h = int(x), int(y), int(w), int(h) @@ -547,4 +555,4 @@ def check_ocr_box(image_source: Union[str, Image.Image], display_img = True, out bb = [get_xywh(item) for item in coord] elif output_bb_format == 'xyxy': bb = [get_xyxy(item) for item in coord] - return (text, bb), goal_filtering \ No newline at end of file + return (text, bb), goal_filtering