Update app.py

app.py

@@ -6,12 +6,11 @@ import os
 import logging
 from pathlib import Path
 from gliner import GLiNER
-from io import BytesIO
 import cv2
 import re
 from PIL import Image
 import traceback
-import tempfile
+import io  # For in-memory file handling
 from difflib import SequenceMatcher
 
 # Configure logging
@@ -25,16 +24,12 @@ os.environ['TRANSFORMERS_CACHE'] = '/tmp/.gliner_models/cache'
 def initialize_models():
     """Initialize models with error handling and retries"""
     try:
-        # Initialize EasyOCR
         logger.info("Initializing EasyOCR...")
         reader = easyocr.Reader(['en', 'ar'], 
                               download_enabled=True,
                               model_storage_directory='/tmp/.easyocr_models')
-        
-        # Initialize GLiNER
         logger.info("Initializing GLiNER...")
         model_path = Path(os.environ['GLINER_HOME']) / 'gliner_large-v2.1'
-        
         if not model_path.exists():
             logger.info("Downloading GLiNER model...")
             model_path.parent.mkdir(parents=True, exist_ok=True)
@@ -42,10 +37,8 @@ def initialize_models():
             model.save_pretrained(str(model_path))
         else:
             model = GLiNER.from_pretrained(str(model_path))
-            
         logger.info("Models initialized successfully")
         return reader, model
-        
     except Exception as e:
         logger.error(f"Model initialization failed: {str(e)}")
         raise
@@ -59,29 +52,31 @@ except Exception as e:
 def clean_extracted_text(text):
     """Clean the extracted text with proper error handling"""
     try:
-        # Preserve Arabic and basic Latin characters along with digits and common punctuation
         cleaned = re.sub(
             r'[^\u0600-\u06FF\u0750-\u077F\u08A0-\u08FFA-Za-z0-9\s@.,-]', 
             '', 
             text
         )
-        # Normalize whitespace
         return re.sub(r'\s+', ' ', cleaned).strip()
     except Exception as e:
         logger.error(f"Text cleaning failed: {traceback.format_exc()}")
-        return text  # Return raw text as fallback
+        return text
 
-def preprocess_image(image):
-    """Image preprocessing with validation"""
+def preprocess_image(image, max_dim=1024):
+    """Image preprocessing with validation and optional resizing"""
     try:
         if not isinstance(image, np.ndarray):
             image = np.array(image)
-            
-        if len(image.shape) == 2:  # Already grayscale
+        # Optional: Resize if the image is too large (keeping aspect ratio)
+        h, w = image.shape[:2]
+        if max(h, w) > max_dim:
+            scaling = max_dim / float(max(h, w))
+            image = cv2.resize(image, (int(w * scaling), int(h * scaling)))
+        # Convert to grayscale if needed
+        if len(image.shape) == 2:
             gray = image
         else:
             gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
-            
         denoised = cv2.medianBlur(gray, 3)
         _, thresh = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
         return thresh
@@ -103,47 +98,35 @@ def clean_and_deduplicate(entities):
     for label, values in entities.items():
         unique = []
         for val in values:
-            # Validate and clean specific fields
             if label.lower() == "email":
                 match = re.search(r'[\w\.-]+@[\w\.-]+', val)
                 val = match.group(0) if match else val
             elif label.lower() == "phone":
                 match = re.search(r'\+?\d[\d\s\-]{7,}\d', val)
                 val = match.group(0) if match else val
-            
-            # Avoid adding duplicates or near-duplicates
             if not any(similar(val, exist) for exist in unique):
                 unique.append(val)
         cleaned_results[label] = unique
     return cleaned_results
 
 def process_single_image(image, threshold=0.3, nested_ner=True, progress=gr.Progress()):
-    """Process single image with detailed error handling and improved entity cleanup"""
+    """Process single image with detailed error handling, optimized I/O, and entity cleanup"""
     try:
-        # Validate input
         if image is None:
             raise ValueError("No image provided")
-            
         progress(0.1, "Validating input...")
         if not isinstance(image, (Image.Image, np.ndarray)):
             raise TypeError(f"Invalid image type: {type(image)}")
-
-        # Preprocessing
         progress(0.2, "Preprocessing image...")
         preprocessed = preprocess_image(image)
-        
-        # OCR
         progress(0.4, "Performing OCR...")
         try:
             ocr_results = reader.readtext(preprocessed, detail=0, paragraph=True)
         except Exception as e:
             logger.error(f"OCR failed: {traceback.format_exc()}")
             raise RuntimeError("OCR processing failed") from e
-            
         raw_text = " ".join(ocr_results)
         clean_text = clean_extracted_text(raw_text)
-        
-        # Entity extraction
         progress(0.6, "Extracting entities...")
         try:
             labels = ["person name", "company name", "job title", "phone", "email", "address"]
@@ -156,38 +139,32 @@ def process_single_image(image, threshold=0.3, nested_ner=True, progress=gr.Prog
         except Exception as e:
             logger.error(f"Entity extraction failed: {traceback.format_exc()}")
             raise RuntimeError("Entity extraction failed") from e
-        
-        # Format raw results into a dictionary by label
         results = {label.title(): [] for label in labels}
         for entity in entities:
             label = entity["label"].title()
             if label in results:
                 results[label].append(entity["text"])
-
-        # Post-process the extracted entities for deduplication and validation
         cleaned_entities = clean_and_deduplicate(results)
-
-        # Create temporary CSV file with final results
-        with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp_file:
-            pd.DataFrame([{
-                k: "; ".join(v) for k, v in cleaned_entities.items()
-            }]).to_csv(tmp_file.name, index=False)
-            csv_path = tmp_file.name
-
+        
+        # Generate CSV output in-memory to reduce disk I/O
+        csv_io = io.BytesIO()
+        pd.DataFrame([{k: "; ".join(v) for k, v in cleaned_entities.items()}]).to_csv(csv_io, index=False)
+        csv_io.seek(0)
+        
         return (
             clean_text,                                  # Text output (str)
             {k: "; ".join(v) for k, v in cleaned_entities.items()},  # JSON output (dict)
-            csv_path,                                    # File path (str)
+            csv_io,                                      # In-memory file (BytesIO)
             ""                                           # Empty error message (str)
         )
         
     except Exception as e:
         logger.error(f"Processing failed: {traceback.format_exc()}")
         return (
-            "",   # Empty text
-            {},   # Empty JSON
-            None, # No file
-            f"Error: {str(e)}\n{traceback.format_exc()}"  # Error details
+            "",   
+            {},    
+            None,  
+            f"Error: {str(e)}\n{traceback.format_exc()}"
         )
 
 # Gradio Interface setup
@@ -207,7 +184,6 @@ with gr.Blocks() as app:
                 error_output = gr.Textbox(label="Error Details", visible=False)
                 csv_download_single = gr.File(label="Download Results")
 
-    # Update click handler to show errors
     submit_single.click(
         fn=process_single_image,
         inputs=[single_image, threshold_single, nested_ner_single],
@@ -223,4 +199,4 @@ app.launch(
     debug=True,
     show_error=True,
     share=False
-)
+)