[feat] llama3.2_vision update

Author: pkarw

app/ocr_strategies/llama_vision.py

@@ -0,0 +1,39 @@
+import base64
+from ocr_strategies.ocr_strategy import OCRStrategy
+import ollama
+import io
+import os
+from pdf2image import convert_from_bytes
+
+class LlamaVisionOCRStrategy(OCRStrategy):
+    """Llama 3.2 Vision OCR Strategy"""
+
+    def extract_text_from_pdf(self, pdf_bytes):
+        # Convert PDF bytes to images
+        images = convert_from_bytes(pdf_bytes)
+        extracted_text = ""
+
+        for i, image in enumerate(images):
+            # Convert image to base64
+            buffered = io.BytesIO()
+            image.save(buffered, format="JPEG")
+            img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+            # Generate text using the Llama 3.2 Vision model
+            try:
+                response = ollama.chat("llama3.2-vision", [{
+                    'content':  os.getenv('LLAMA_VISION_PROMPT', "You are OCR. Convert image to markdown."),
+                    'images': [img_str]
+                }], stream=True)
+                for chunk in response:
+                    extracted_text += chunk['message']['content']
+
+            except ollama.ResponseError as e:
+                print('Error:', e.error)
+                raise Exception("Failed to generate text with Llama 3.2 Vision model")
+
+            print(response) 
+            #page_text = response.get("response", "")
+            #extracted_text += f"--- Page {i + 1} ---\n{page_text}\n"
+
+        return extracted_text

app/tasks.py

@@ -2,14 +2,16 @@
 from celery_config import celery
 from ocr_strategies.marker import MarkerOCRStrategy
 from ocr_strategies.tesseract import TesseractOCRStrategy
+from ocr_strategies.llama_vision import LlamaVisionOCRStrategy
 import redis
 import os
 import ollama
 from storage_manager import StorageManager
 
 OCR_STRATEGIES = {
     'marker': MarkerOCRStrategy(),
-    'tesseract': TesseractOCRStrategy()
+    'tesseract': TesseractOCRStrategy(),
+    'llama_vision': LlamaVisionOCRStrategy()
 }
 
 # Connect to Redis
@@ -23,17 +25,17 @@ def ocr_task(self, pdf_bytes, strategy_name, pdf_filename, pdf_hash, ocr_cache,
     """
     start_time = time.time()
     if strategy_name not in OCR_STRATEGIES:
-        raise ValueError(f"Unknown strategy '{strategy_name}'. Available: marker, tesseract")
+        raise ValueError(f"Unknown strategy '{strategy_name}'. Available: marker, tesseract, llama_vision")
 
     ocr_strategy = OCR_STRATEGIES[strategy_name]
     self.update_state(state='PROGRESS', status="File uploaded successfully", meta={'progress': 10})  # Example progress update
 
     extracted_text = None
-    if ocr_cache:
-        cached_result = redis_client.get(pdf_hash)
-        if cached_result:
-            # Return cached result if available
-            extracted_text = cached_result.decode('utf-8')
+    # if ocr_cache:
+    #     cached_result = redis_client.get(pdf_hash)
+    #     if cached_result:
+    #         # Return cached result if available
+    #         extracted_text = cached_result.decode('utf-8')
 
     if extracted_text is None:
         print("Extracting text from PDF...")