Upload OmniParser/util/omniparser.py with huggingface_hub

OmniParser/util/omniparser.py

@@ -0,0 +1,60 @@
+from util.utils import get_som_labeled_img, get_caption_model_processor, get_yolo_model, check_ocr_box
+import torch
+from PIL import Image
+import io
+import base64
+from typing import Dict
+import os
+from pathlib import Path
+
+class Omniparser(object):
+    def __init__(self, config: Dict):
+        self.config = config
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+        self.som_model = get_yolo_model(model_path=config['som_model_path'])
+        # Disable caption generation for speed on CPU
+        self.caption_model_processor = None
+        
+        # Setup for saving cropped images
+        self.save_cropped_images = config.get('save_cropped_images', False)
+        self.cropped_images_dir = config.get('cropped_images_dir', 'cropped_images')
+        
+        if self.save_cropped_images:
+            Path(self.cropped_images_dir).mkdir(parents=True, exist_ok=True)
+            print(f'Cropped images will be saved to: {self.cropped_images_dir}')
+        
+        print('Omniparser initialized!!!')
+
+    def parse(self, image_base64: str):
+        try:
+            print("[1] Decoding image...")
+            image_bytes = base64.b64decode(image_base64)
+            image = Image.open(io.BytesIO(image_bytes))
+            print(f'[1] Image decoded, size: {image.size}')
+            
+            print("[2] Calculating overlay config...")
+            box_overlay_ratio = max(image.size) / 3200
+            draw_bbox_config = {
+                'text_scale': 0.8 * box_overlay_ratio,
+                'text_thickness': max(int(2 * box_overlay_ratio), 1),
+                'text_padding': max(int(3 * box_overlay_ratio), 1),
+                'thickness': max(int(3 * box_overlay_ratio), 1),
+            }
+            print("[2] Config ready")
+
+            print("[3] Checking OCR boxes...")
+            (text, ocr_bbox), _ = check_ocr_box(image, display_img=False, output_bb_format='xyxy', easyocr_args={'text_threshold': 0.8}, use_paddleocr=False)
+            print(f"[3] OCR complete: {len(text)} text boxes, {len(ocr_bbox) if ocr_bbox else 0} OCR bboxes")
+            
+            print("[4] Getting SOM labeled image...")
+            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, save_cropped_images=self.save_cropped_images, cropped_images_dir=self.cropped_images_dir)
+            print("[4] SOM complete")
+
+            print("[5] Returning results...")
+            return dino_labled_img, parsed_content_list
+        except Exception as e:
+            import traceback
+            print(f"[ERROR] Parse failed: {str(e)}")
+            print(traceback.format_exc())
+            raise