DocuLens/backend/utils.py

# backend/utils.py
import os
import time
import pytesseract
from PIL import Image
from pdf2image import convert_from_path
import shutil
import asyncio
import uuid
import subprocess
import base64
import re
from models import get_prompt

# These imports were only used by TxtToPdf functionality and have been removed

UPLOAD_FOLDER = "uploads"
OUTPUT_FOLDER = "outputs"

def encode_image(file_path):
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def run_tesseract(file_path):
    """
    Si el archivo es un PDF, se realiza OCR página a página añadiendo al principio de cada una
    la cabecera con el formato [Page 0001], [Page 0002], etc.
    Para otros formatos se realiza OCR normal.
    """
    extracted_text = ""
    if file_path.lower().endswith(".pdf"):
        try:
            pages = convert_from_path(file_path)
            for i, page in enumerate(pages, start=1):
                page_text = pytesseract.image_to_string(page, lang='eng')
                extracted_text += f"[Page {i:04d}]\n{page_text}\n\n"
        except Exception as e:
            extracted_text = f"❌ Error processing PDF: {str(e)}"
    else:
        image = Image.open(file_path)
        extracted_text = pytesseract.image_to_string(image, lang='eng')
    return extracted_text

def call_api_correction(api, model, text, prompt_key="ocr_correction"):
    prompt = get_prompt(prompt_key) + text
    try:
        import google.generativeai as genai
    except ImportError:
        raise ImportError("Please install 'google-generativeai' to use Gemini.")
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    
    # Get the API key
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return "❌ GEMINI_API_KEY not found in environment variables"
    
    # Configure API
    genai.configure(api_key=api_key)
    
    # Create model instance
    text_model = genai.GenerativeModel(model)
    
    # Generate content
    response = text_model.generate_content(prompt)
    
    # Get text response
    return response.text

def call_api_ocr(api, model, file_path, prompt_key="ocr"):
    """
    OCR processing using Gemini's vision capabilities.
    Implements image optimization and retries with exponential backoff for handling timeouts.
    """
    try:
        import sys
        import traceback
        import mimetypes
        import base64
        import time
        from PIL import Image
        import google.generativeai as genai
    except ImportError:
        raise ImportError("Please install 'google-generativeai' to use Gemini.")
    
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
    MAX_RETRIES = 3
    RETRY_DELAY = 2  # Initial delay in seconds, will increase with each retry
    
    try:
        # Debug information
        print(f"DEBUG OCR: File path: {file_path}")
        print(f"DEBUG OCR: File exists: {os.path.exists(file_path)}")
        print(f"DEBUG OCR: File size: {os.path.getsize(file_path) if os.path.exists(file_path) else 'N/A'}")
        
        # Get the API key
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            error_msg = "GEMINI_API_KEY not found in environment variables"
            print(f"DEBUG OCR ERROR: {error_msg}")
            return f"❌ {error_msg}"
            
        # Configure API
        print(f"DEBUG OCR: Configuring API with key: {api_key[:4]}...{api_key[-4:] if len(api_key) > 8 else ''}")
        genai.configure(api_key=api_key)
        
        # Open and optimize the image
        try:
            # Open the image
            img = Image.open(file_path)
            
            # Check image size - resize if too large (max 1600px on any side)
            max_size = 1600
            original_size = img.size
            if img.width > max_size or img.height > max_size:
                # Calculate new size preserving aspect ratio
                if img.width > img.height:
                    new_width = max_size
                    new_height = int(img.height * (max_size / img.width))
                else:
                    new_height = max_size
                    new_width = int(img.width * (max_size / img.height))
                
                img = img.resize((new_width, new_height), Image.LANCZOS)
                print(f"DEBUG OCR: Resized image from {original_size} to {img.size}")
            
            # Convert to RGB if needed (e.g., for PNG with transparency)
            if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
                background = Image.new('RGB', img.size, (255, 255, 255))
                if img.mode == 'P':
                    img = img.convert('RGBA')
                background.paste(img, (0, 0), img)
                img = background
                print("DEBUG OCR: Converted image to RGB")
            elif img.mode != 'RGB':
                img = img.convert('RGB')
                print("DEBUG OCR: Converted image to RGB")
            
            # Save to BytesIO and encode
            import io
            buffer = io.BytesIO()
            img.save(buffer, format="JPEG", quality=85, optimize=True)
            buffer.seek(0)
            image_bytes = buffer.read()
            file_base64 = base64.b64encode(image_bytes).decode('utf-8')
            print(f"DEBUG OCR: Successfully processed and encoded image ({len(file_base64)} chars)")
        except Exception as img_error:
            error_msg = f"Error processing image: {str(img_error)}"
            print(f"DEBUG OCR ERROR: {error_msg}")
            print(traceback.format_exc())
            return f"❌ {error_msg}"
            
        # Get prompt
        prompt = get_prompt(prompt_key)
        print(f"DEBUG OCR: Using prompt: {prompt[:50]}...")
        
        # Create model instance
        vision_model = genai.GenerativeModel(model)
        print(f"DEBUG OCR: Created model instance for model: {model}")
        
        # Create image data in the format required by the API
        image_data = {
            "mime_type": "image/jpeg",  # Always using JPEG for consistency
            "data": file_base64
        }
        
        print(f"DEBUG OCR: Image prepared for API. Generating content with retries...")
        
        # Initialize variables for retry logic
        retry_count = 0
        current_delay = RETRY_DELAY
        last_error = None
        
        # Retry loop
        while retry_count < MAX_RETRIES:
            try:
                if retry_count > 0:
                    print(f"DEBUG OCR: Retry {retry_count}/{MAX_RETRIES} after {current_delay}s delay")
                
                # Generate content with the image and the prompt
                response = vision_model.generate_content([
                    {"inline_data": image_data},
                    prompt
                ], generation_config={
                    "temperature": 0.2,
                    "top_p": 0.95,
                    "top_k": 40,
                    "max_output_tokens": 1024,  # Limit output size to avoid timeouts
                })
                
                print(f"DEBUG OCR: Response received successfully")
                
                # Return the text response
                return response.text
                
            except Exception as api_error:
                last_error = api_error
                error_str = str(api_error)
                print(f"DEBUG OCR ERROR on attempt {retry_count+1}: {error_str}")
                
                # Check if this is a timeout or rate limit error (worth retrying)
                if "Deadline Exceeded" in error_str or "429" in error_str or "quota" in error_str.lower() or "timeout" in error_str.lower():
                    retry_count += 1
                    if retry_count < MAX_RETRIES:
                        print(f"DEBUG OCR: Timeout or rate limit error, will retry in {current_delay}s")
                        time.sleep(current_delay)
                        current_delay *= 2  # Exponential backoff
                    else:
                        print("DEBUG OCR: Max retries reached")
                else:
                    # Non-retryable error
                    print("DEBUG OCR: Non-retryable error, giving up")
                    break
        
        # If we get here, all retries failed
        error_detail = f"Error after {MAX_RETRIES} attempts: {str(last_error)}"
        print(f"DEBUG OCR ERROR: {error_detail}")
        return f"❌ Error in OCR: {error_detail}"
        
    except Exception as e:
        error_detail = f"Error: {str(e)}\n"
        error_detail += f"Exception type: {type(e).__name__}\n"
        error_detail += f"Traceback: {traceback.format_exc()}"
        print(f"DEBUG OCR ERROR: {error_detail}")
        
        return f"❌ Error in OCR: {str(e)}"

def describe_image(file_path, api, model, prompt_key="image_desc", update_progress=None, is_cancelled=None):
    """
    Process an image file and generate a description using the specified AI model.
    Uses direct base64 encoding for Gemini's vision capabilities.
    Implements retries with exponential backoff for handling timeouts.
    """
    import sys
    import traceback
    import mimetypes
    import base64
    import time
    from PIL import Image
    
    print(f"DEBUG IMAGE_DESC: Starting image description function for file: {file_path}")
    print(f"DEBUG IMAGE_DESC: Using API: {api}, Model: {model}, Prompt key: {prompt_key}")
    
    MAX_RETRIES = 3
    RETRY_DELAY = 2  # Initial delay in seconds, will increase with each retry
    
    try:
        if update_progress:
            update_progress(20, "⏳ Loading image file...")
        
        if is_cancelled and is_cancelled():
            return "Process cancelled."
        
        # Check if file exists
        if not os.path.exists(file_path):
            error_msg = f"File not found: {file_path}"
            print(f"DEBUG IMAGE_DESC: {error_msg}")
            if update_progress:
                update_progress(0, f"❌ {error_msg}")
            return f"❌ {error_msg}"
            
        # Get MIME type of image
        mime_type = mimetypes.guess_type(file_path)[0] or 'image/png'
        print(f"DEBUG IMAGE_DESC: File path: {file_path}")
        print(f"DEBUG IMAGE_DESC: MIME type: {mime_type}")
        
        # Read the image, resize if needed to reduce size, and encode it to base64
        if update_progress:
            update_progress(40, "🔍 Processing image...")
            
        try:
            # Open and optimize the image
            img = Image.open(file_path)
            
            # Check image size - resize if too large (max 1600px on any side)
            max_size = 1600
            original_size = img.size
            if img.width > max_size or img.height > max_size:
                # Calculate new size preserving aspect ratio
                if img.width > img.height:
                    new_width = max_size
                    new_height = int(img.height * (max_size / img.width))
                else:
                    new_height = max_size
                    new_width = int(img.width * (max_size / img.height))
                
                img = img.resize((new_width, new_height), Image.LANCZOS)
                print(f"DEBUG IMAGE_DESC: Resized image from {original_size} to {img.size}")
            
            # Convert to RGB if needed (e.g., for PNG with transparency)
            if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
                # Create a white background image
                background = Image.new('RGB', img.size, (255, 255, 255))
                # Paste the image on the background (handles transparency)
                if img.mode == 'P':
                    img = img.convert('RGBA')
                background.paste(img, (0, 0), img)
                img = background
                print("DEBUG IMAGE_DESC: Converted image to RGB")
            elif img.mode != 'RGB':
                img = img.convert('RGB')
                print("DEBUG IMAGE_DESC: Converted image to RGB")
            
            # Save to BytesIO and encode
            import io
            buffer = io.BytesIO()
            img.save(buffer, format="JPEG", quality=85, optimize=True)
            buffer.seek(0)
            image_bytes = buffer.read()
            image_base64 = base64.b64encode(image_bytes).decode('utf-8')
            print(f"DEBUG IMAGE_DESC: Successfully processed and encoded image ({len(image_base64)} chars)")
        except Exception as e:
            error_msg = f"Error processing image file: {str(e)}"
            print(f"DEBUG IMAGE_DESC ERROR: {error_msg}")
            print(traceback.format_exc())
            if update_progress:
                update_progress(0, f"❌ {error_msg}")
            return f"❌ {error_msg}"
        
        # Import Gemini libraries - using the correct package name
        try:
            print("DEBUG IMAGE_DESC: Attempting to import google-generativeai library")
            import google.generativeai as genai
            print("DEBUG IMAGE_DESC: Successfully imported google.generativeai library")
        except ImportError as ie:
            error_msg = f"Failed to import Google's generativeai library: {str(ie)}. Please install it with: pip install google-generativeai"
            print(f"DEBUG IMAGE_DESC ERROR: {error_msg}")
            if update_progress:
                update_progress(0, f"❌ {error_msg}")
            return f"❌ {error_msg}"
        
        # Configure API 
        api_key = os.environ.get("GEMINI_API_KEY")
        if not api_key:
            error_msg = "GEMINI_API_KEY not found in environment variables"
            print(f"DEBUG IMAGE_DESC ERROR: {error_msg}")
            if update_progress:
                update_progress(0, f"❌ {error_msg}")
            return f"❌ {error_msg}"
        
        print(f"DEBUG IMAGE_DESC: Configuring API with key: {api_key[:4]}...{api_key[-4:] if len(api_key) > 8 else ''}")
        genai.configure(api_key=api_key)
        
        # Get prompt from configuration
        prompt = get_prompt(prompt_key)
        print(f"DEBUG IMAGE_DESC: Using prompt: {prompt[:50]}...")
        
        if update_progress:
            update_progress(60, "🤖 Generating description...")
        
        # Generation parameters - reduce complexity for better performance
        generation_config = {
            "temperature": 0.2,  # Lower temperature for more deterministic output
            "top_p": 0.95,
            "top_k": 40,
            "max_output_tokens": 800,  # Limit output size
        }
        
        print(f"DEBUG IMAGE_DESC: Creating model instance for model: {model}")
        # Create model instance
        vision_model = genai.GenerativeModel(model)
        
        # Format image for the API
        image_data = {
            "mime_type": "image/jpeg",  # Always using JPEG for consistency
            "data": image_base64
        }
        
        print(f"DEBUG IMAGE_DESC: Image prepared for API. Generating content with retries...")
        
        # Initialize variables for retry logic
        retry_count = 0
        current_delay = RETRY_DELAY
        last_error = None
        
        # Retry loop
        while retry_count < MAX_RETRIES:
            try:
                if retry_count > 0:
                    print(f"DEBUG IMAGE_DESC: Retry {retry_count}/{MAX_RETRIES} after {current_delay}s delay")
                    if update_progress:
                        update_progress(60, f"🔄 Retry {retry_count}/{MAX_RETRIES}...")
                
                # Generate content using the correct format for v0.3.1
                response = vision_model.generate_content([
                    {"inline_data": image_data},
                    prompt
                ], generation_config=generation_config)
                
                print(f"DEBUG IMAGE_DESC: Successfully received response from Gemini API")
                
                if update_progress:
                    update_progress(100, "✅ Description generated")
                
                # Get text response 
                print(f"DEBUG IMAGE_DESC: Returning text response of length: {len(response.text) if response.text else 0}")
                return response.text
                
            except Exception as api_error:
                last_error = api_error
                error_str = str(api_error)
                print(f"DEBUG IMAGE_DESC ERROR on attempt {retry_count+1}: {error_str}")
                
                # Check if this is a timeout or rate limit error (worth retrying)
                if "Deadline Exceeded" in error_str or "429" in error_str or "quota" in error_str.lower() or "timeout" in error_str.lower():
                    retry_count += 1
                    if retry_count < MAX_RETRIES:
                        print(f"DEBUG IMAGE_DESC: Timeout or rate limit error, will retry in {current_delay}s")
                        time.sleep(current_delay)
                        current_delay *= 2  # Exponential backoff
                    else:
                        print("DEBUG IMAGE_DESC: Max retries reached")
                else:
                    # Non-retryable error
                    print("DEBUG IMAGE_DESC: Non-retryable error, giving up")
                    break
        
        # If we get here, all retries failed
        error_msg = f"Error generating content after {MAX_RETRIES} attempts: {str(last_error)}"
        print(f"DEBUG IMAGE_DESC ERROR: {error_msg}")
        if update_progress:
            update_progress(0, f"❌ {error_msg}")
        return f"❌ {error_msg}"
            
    except Exception as e:
        error_msg = f"Unexpected error: {str(e)}"
        print(f"DEBUG IMAGE_DESC ERROR: {error_msg}")
        print(traceback.format_exc())
        if update_progress:
            update_progress(0, f"❌ {error_msg}")
        return f"❌ {error_msg}"

def call_api_translation(api, model, text, target_language, prompt_key="translation"):
    prompt_template = get_prompt(prompt_key)
    prompt = prompt_template.format(target_language=target_language) + text
    try:
        import google.generativeai as genai
    except ImportError:
        raise ImportError("Please install 'google-generativeai' to use Gemini.")
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
    # Get the API key
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return "❌ GEMINI_API_KEY not found in environment variables"
    
    # Configure API
    genai.configure(api_key=api_key)
    
    # Create model instance
    text_model = genai.GenerativeModel(model)
    
    # Generate content - simpler format for v0.3.1
    response = text_model.generate_content(prompt, generation_config={
        "temperature": 0.2,
        "top_p": 0.95,
        "top_k": 40
    })
    
    # Get text response
    return response.text

def embed_ocr_in_pdf(input_pdf, output_pdf):
    try:
        subprocess.run(["ocrmypdf", input_pdf, output_pdf], check=True)
        return True
    except Exception as e:
        return False

def ocr_file_by_pages(file_path, api, model, prompt_key, update_progress, is_cancelled):
    final_text = ""
    if file_path.lower().endswith(".pdf"):
        try:
            pages = convert_from_path(file_path)
        except Exception as e:
            return f"❌ Error processing PDF: {str(e)}"
        total = len(pages)
        for i, page in enumerate(pages, start=1):
            if is_cancelled():
                update_progress(0, "⏹️ Cancelled")
                return "Process cancelled."
            temp_filename = os.path.join(OUTPUT_FOLDER, f"temp_page_{uuid.uuid4().hex}.png")
            page.save(temp_filename, "PNG")
            page_text = call_api_ocr(api, model, temp_filename, prompt_key)
            final_text += f"[Page {i:04d}]\n{page_text}\n\n"
            os.remove(temp_filename)
            progress = int((i / total) * 100)
            update_progress(progress, f"📄 Processed page {i} of {total}.")
            time.sleep(1)
        return final_text
    else:
        return call_api_ocr(api, model, file_path, prompt_key)

def translate_file_by_pages(file_path, api, model, target_language, prompt_key, update_progress, is_cancelled):
    final_translation = ""
    if file_path.lower().endswith(".pdf"):
        try:
            pages = convert_from_path(file_path)
        except Exception as e:
            return f"❌ Error processing PDF: {str(e)}"
        total = len(pages)
        for i, page in enumerate(pages, start=1):
            if is_cancelled():
                update_progress(0, "⏹️ Cancelled")
                return "Process cancelled."
            temp_filename = os.path.join(OUTPUT_FOLDER, f"temp_page_{uuid.uuid4().hex}.png")
            page.save(temp_filename, "PNG")
            page_text = pytesseract.image_to_string(page, lang='eng')
            translated_page = call_api_translation(api, model, page_text, target_language, prompt_key)
            final_translation += f"[Page {i:04d}]\n{translated_page}\n\n"
            os.remove(temp_filename)
            progress = int((i / total) * 100)
            update_progress(progress, f"📄 Processed page {i} of {total}.")
            time.sleep(1)
        return final_translation
    elif file_path.lower().endswith(".txt"):
        with open(file_path, "r", encoding="utf-8") as f:
            text = f.read()
        translated = call_api_translation(api, model, text, target_language, prompt_key)
        update_progress(100, "🎉 Process completed")
        return f"[Page {1:04d}]\n{translated}"
    else:
        return "Unsupported file type for translation."

def process_file(file_path, api, model, mode, prompt_key, update_progress, is_cancelled):
    """
    NOTA: Ahora solo usamos [Page XXXX] al inicio de cada página (por Tesseract).
    """
    if is_cancelled():
        update_progress(0, "⏹️ Cancelled")
        return "Process cancelled."
    update_progress(25, "📤 File uploaded.")
    base_name = os.path.splitext(os.path.basename(file_path))[0]

    if mode == "OCR":
        processed_text = run_tesseract(file_path)
        if is_cancelled():
            update_progress(25, "⏹️ Cancelled")
            return "Process cancelled."
        update_progress(50, "✅ Tesseract OCR completed.")
        pdf_output = os.path.join(OUTPUT_FOLDER, base_name + "_ocr.pdf")
        if file_path.lower().endswith(".pdf"):
            if embed_ocr_in_pdf(file_path, pdf_output):
                update_progress(95, "📄 OCR embedded into PDF.")
            else:
                shutil.copy(file_path, pdf_output)
                update_progress(95, "⚠️ Failed to embed OCR; original PDF copied.")
        else:
            shutil.copy(file_path, pdf_output)

    elif mode == "OCR + AI":
        processed_text = run_tesseract(file_path)
        if is_cancelled():
            update_progress(25, "⏹️ Cancelled")
            return "Process cancelled."
        update_progress(50, "✅ Tesseract OCR completed.")
        processed_text = call_api_correction(api, model, processed_text, prompt_key)
        pdf_output = os.path.join(OUTPUT_FOLDER, base_name + "_ocr.pdf")
        if file_path.lower().endswith(".pdf"):
            if embed_ocr_in_pdf(file_path, pdf_output):
                update_progress(95, "📄 OCR embedded into PDF.")
            else:
                shutil.copy(file_path, pdf_output)
                update_progress(95, "⚠️ Failed to embed OCR; original PDF copied.")
        else:
            shutil.copy(file_path, pdf_output)

    elif mode == "AI":
        update_progress(25, "📂 File ready for full AI processing.")
        if file_path.lower().endswith(".pdf"):
            processed_text = ocr_file_by_pages(file_path, api, model, prompt_key, update_progress, is_cancelled)
        else:
            processed_text = call_api_ocr(api, model, file_path, prompt_key)

    else:
        processed_text = "Unrecognized processing mode."
        update_progress(25, "❌ Error: Unrecognized mode.")

    if is_cancelled():
        update_progress(75, "⏹️ Cancelled")
        return "Process cancelled."

    update_progress(75, "🤖 API processing completed.")

    # Guardar el TXT final
    txt_file = os.path.join(OUTPUT_FOLDER, base_name + ".txt")
    with open(txt_file, "w", encoding="utf-8") as f:
        f.write(processed_text)

    update_progress(100, "🎉 Process completed")
    return processed_text