image processing resposeDocument formate compatible endpoint - Adithya S k

Author: adithya-s-k

omniparse/__init__.py

@@ -23,7 +23,7 @@ def load_omnimodel(load_documents: bool, load_media: bool, load_web: bool):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if load_documents:
         print("[LOG] ✅ Loading OCR Model")
-        shared_state.model_list = load_all_models()
+        # shared_state.model_list = load_all_models()
         print("[LOG] ✅ Loading Vision Model")
         # if device == "cuda":
         shared_state.vision_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device)

omniparse/image/__init__.py

@@ -6,8 +6,9 @@
 from PIL import Image
 # from omniparse.document.parse import parse_single_image
 from omniparse.documents.parse import parse_single_pdf
-from omniparse.image.process import pre_process_image
+from omniparse.image.process import process_image_task
 from omniparse.utils import encode_images
+from omniparse.models import responseDocument
 
 def parse_image(input_data, model_state) -> dict:
     temp_files = []
@@ -44,17 +45,22 @@ def parse_image(input_data, model_state) -> dict:
 
         # Parse the PDF file
         full_text, images, out_meta = parse_single_pdf(temp_pdf_path, model_state.model_list)
-        images = encode_images(images)
+        
+        parse_image_result = responseDocument(
+            text=full_text,
+            metadata=out_meta
+        )
+        encode_images(images,parse_image_result)
 
-        return {"message": "Document parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images}
+        return parse_image_result
 
     finally:
         # Clean up the temporary files
         for file_path in temp_files:
             if os.path.exists(file_path):
                 os.remove(file_path)
 
-def process_image(input_data, task, model_state) -> dict:
+def process_image(input_data, task, model_state) -> responseDocument:
     try:
         temp_files = []
 
@@ -76,9 +82,9 @@ def process_image(input_data, task, model_state) -> dict:
         image_data = Image.open(temp_file_path).convert("RGB")
 
         # Process the image using your function (e.g., process_image)
-        results = pre_process_image(image_data, task, vision_model = model_state.vision_model, vision_processor = model_state.vision_processor)
+        image_process_results : responseDocument = process_image_task(image_data, task, model_state)
 
-        return {"results": results}
+        return image_process_results
 
     finally:
         # Clean up the temporary files

omniparse/image/parse.py

@@ -1,9 +1,7 @@
 import warnings
-warnings.filterwarnings("ignore", category=UserWarning) # Filter torch pytree user warnings
-
+from typing import List, Dict, Tuple, Optional
 import pypdfium2 as pdfium # Needs to be at the top to avoid warnings
 from PIL import Image
-
 from omniparse.documents.utils import flush_cuda_memory
 from omniparse.documents.tables.table import format_tables
 from omniparse.documents.debug.data import dump_bbox_debug_data
@@ -25,14 +23,14 @@
 from omniparse.documents.cleaners.text import cleanup_text
 from omniparse.documents.images.extract import extract_images
 from omniparse.documents.images.save import images_to_dict
-
-from typing import List, Dict, Tuple, Optional
 from omniparse.documents.settings import settings
 
+warnings.filterwarnings("ignore", category=UserWarning) # Filter torch pytree user warnings
+
 
 def parse_single_image(
         image: Image.Image,
-        model_lst: List,
+        model_list: List,
         metadata: Optional[Dict] = None,
         langs: Optional[List[str]] = None,
         batch_multiplier: int = 1
@@ -54,10 +52,10 @@ def parse_single_image(
         "languages": langs,
     }
 
-    texify_model, layout_model, order_model, edit_model, detection_model, ocr_model = model_lst
+    texify_model, layout_model, order_model, edit_model, detection_model, ocr_model = model_list
 
     # Identify text lines on pages
-    text_line_prediction = surya_detection(image, detection_model, batch_multiplier=batch_multiplier)
+    surya_detection(image, detection_model, batch_multiplier=batch_multiplier)
     flush_cuda_memory()
 
     # OCR pages as needed

omniparse/image/process.py

@@ -1,53 +1,144 @@
+from typing import Dict, Any, Union
+from PIL import Image as PILImage
+import base64
+from io import BytesIO
+import copy
+from omniparse.image.utils import plot_bbox, fig_to_pil,draw_polygons,draw_ocr_bboxes
+from omniparse.models import responseDocument
 
-def pre_process_image(image, task_prompt, vision_model, vision_processor):
-    # :Convert binary image data to PIL Image
-    # image = Image.fromarray(image) 
+def process_image_task(image_data: Union[str, bytes, PILImage.Image], task_prompt: str,  model_state) -> Dict[str, Any]:
+    # Convert image_data if it's in bytes
+    if isinstance(image_data, bytes):
+        pil_image = PILImage.open(BytesIO(image_data))
+    elif isinstance(image_data, str):
+        try:
+            image_bytes = base64.b64decode(image_data)
+            pil_image = PILImage.open(BytesIO(image_bytes))
+        except Exception as e:
+            raise ValueError(f"Failed to decode base64 image: {str(e)}")
+    elif isinstance(image_data, PILImage.Image):
+        pil_image = image_data
+    else:
+        raise ValueError("Unsupported image_data type. Should be either string (file path), bytes (binary image data), or PIL.Image instance.")
+
+    # Process based on task_prompt
     if task_prompt == 'Caption':
-        task_prompt = '<CAPTION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<CAPTION>'
     elif task_prompt == 'Detailed Caption':
-        task_prompt = '<DETAILED_CAPTION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<DETAILED_CAPTION>'
     elif task_prompt == 'More Detailed Caption':
-        task_prompt = '<MORE_DETAILED_CAPTION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<MORE_DETAILED_CAPTION>'
+    elif task_prompt == 'Caption + Grounding':
+        task_prompt_model = '<CAPTION>'
+    elif task_prompt == 'Detailed Caption + Grounding':
+        task_prompt_model = '<DETAILED_CAPTION>'
+    elif task_prompt == 'More Detailed Caption + Grounding':
+        task_prompt_model = '<MORE_DETAILED_CAPTION>'
     elif task_prompt == 'Object Detection':
-        task_prompt = '<OD>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<OD>'
     elif task_prompt == 'Dense Region Caption':
-        task_prompt = '<DENSE_REGION_CAPTION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<DENSE_REGION_CAPTION>'
     elif task_prompt == 'Region Proposal':
-        task_prompt = '<REGION_PROPOSAL>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<REGION_PROPOSAL>'
     elif task_prompt == 'Caption to Phrase Grounding':
-        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<CAPTION_TO_PHRASE_GROUNDING>'
     elif task_prompt == 'Referring Expression Segmentation':
-        task_prompt = '<REFERRING_EXPRESSION_SEGMENTATION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<REFERRING_EXPRESSION_SEGMENTATION>'
     elif task_prompt == 'Region to Segmentation':
-        task_prompt = '<REGION_TO_SEGMENTATION>'
-        results = run_example(task_prompt, image,vision_model, vision_processor)
+        task_prompt_model = '<REGION_TO_SEGMENTATION>'
     elif task_prompt == 'Open Vocabulary Detection':
-        task_prompt = '<OPEN_VOCABULARY_DETECTION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<OPEN_VOCABULARY_DETECTION>'
     elif task_prompt == 'Region to Category':
-        task_prompt = '<REGION_TO_CATEGORY>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<REGION_TO_CATEGORY>'
     elif task_prompt == 'Region to Description':
-        task_prompt = '<REGION_TO_DESCRIPTION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<REGION_TO_DESCRIPTION>'
     elif task_prompt == 'OCR':
-        task_prompt = '<OCR>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<OCR>'
     elif task_prompt == 'OCR with Region':
-        task_prompt = '<OCR_WITH_REGION>'
-        results = run_example(task_prompt, image, vision_model, vision_processor)
+        task_prompt_model = '<OCR_WITH_REGION>'
     else:
-        return {"error": "Invalid task prompt"}
+        raise ValueError("Invalid task prompt")
+
+    results, processed_image = pre_process_image(pil_image, task_prompt_model, model_state.vision_model, model_state.vision_processor)
+    # Update responseDocument fields based on the results
+    process_image_result = responseDocument(
+        text = str(results)
+    )
+
+    if processed_image is not None:
+        process_image_result.add_image(f"{task_prompt}", processed_image)
+
+    return process_image_result
 
-    return results
+# Your pre_process_image function with some adjustments
+def pre_process_image(image, task_prompt, vision_model, vision_processor):
+    if task_prompt == '<CAPTION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        return results, None
+    elif task_prompt == '<DETAILED_CAPTION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        return results, None
+    elif task_prompt == '<MORE_DETAILED_CAPTION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        return results, None
+    elif task_prompt == '<CAPTION_TO_PHRASE_GROUNDING>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<DETAILED_CAPTION + GROUNDING>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<MORE_DETAILED_CAPTION + GROUNDING>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<OD>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<DENSE_REGION_CAPTION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<REGION_PROPOSAL>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<CAPTION_TO_PHRASE_GROUNDING>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<REFERRING_EXPRESSION_SEGMENTATION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results[task_prompt], fill_mask=True)
+        return results, output_image
+    elif task_prompt == '<REGION_TO_SEGMENTATION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results[task_prompt], fill_mask=True)
+        return results, output_image
+    elif task_prompt == '<OPEN_VOCABULARY_DETECTION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        fig = plot_bbox(image, results[task_prompt])
+        return results, fig_to_pil(fig)
+    elif task_prompt == '<REGION_TO_CATEGORY>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        return results, None
+    elif task_prompt == '<REGION_TO_DESCRIPTION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        return results, None
+    elif task_prompt == '<OCR>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        return results, None
+    elif task_prompt == '<OCR_WITH_REGION>':
+        results = run_example(task_prompt, image, vision_model, vision_processor)
+        output_image = copy.deepcopy(image)
+        output_image = draw_ocr_bboxes(output_image, results[task_prompt])
+        return results, output_image
+    else:
+        raise ValueError("Invalid task prompt")
 
 def run_example(task_prompt, image, vision_model, vision_processor):
     # if text_input is None:

omniparse/image/router.py

@@ -2,6 +2,7 @@
 from fastapi.responses import JSONResponse
 from omniparse import get_shared_state
 from omniparse.image import parse_image, process_image
+from omniparse.models import responseDocument
 
 image_router = APIRouter()
 model_state = get_shared_state()
@@ -20,8 +21,8 @@ async def parse_image_endpoint(file: UploadFile = File(...)):
 async def process_image_route(image: UploadFile = File(...), task: str = Form(...)):
     try:
         file_bytes = await image.read()
-        result = process_image(file_bytes, task, model_state)
-        return JSONResponse(content=result)
+        result : responseDocument = process_image(file_bytes, task, model_state)
+        return JSONResponse(content=result.model_dump())
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

omniparse/image/utils.py

@@ -1,3 +1,4 @@
+import io
 import random
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
@@ -66,3 +67,8 @@ def draw_ocr_bboxes(image, prediction):
     return image
 
 
+def fig_to_pil(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png')
+    buf.seek(0)
+    return Image.open(buf)

omniparse/utils.py

@@ -1,23 +1,23 @@
 import base64
 import os
 from art import text2art
+from omniparse.models import responseDocument
 
-def encode_images(images):
-    image_data = {}
+def encode_images(images, inputDocument:responseDocument):
+    image_data = []
     for i, (filename, image) in enumerate(images.items()):
         # print(f"Processing image {filename}")
-        
         # Save image as PNG
         image.save(filename, "PNG")
-
         # Read the saved image file as bytes
         with open(filename, "rb") as f:
             image_bytes = f.read()
-
         # Convert image to base64
         image_base64 = base64.b64encode(image_bytes).decode('utf-8')
         image_data[f'{filename}'] = image_base64
-
+        
+        inputDocument.add_image(image_name=filename,image_data=image_base64)
+        
         # Remove the temporary image file
         os.remove(filename)
     return image_data