Update app.py

app.py

@@ -1,20 +1,18 @@
 import os
-import tempfile
-import logging
+import cv2
+import numpy as np
+from PIL import Image
+import pytesseract
 import gradio as gr
-import PyPDF2
 from pdf2image import convert_from_path
-import pytesseract
-from PIL import Image
-import docx
+import PyPDF2
 from llama_index.core import VectorStoreIndex, Document
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.llms.openai import OpenAI
 from llama_index.core import get_response_synthesizer
 from dotenv import load_dotenv
 from sentence_transformers import SentenceTransformer, util
-from sklearn.metrics.pairwise import cosine_similarity
-import numpy as np
+import logging
 
 # Set up logging configuration
 logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)s | %(message)s')
@@ -22,79 +20,63 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)s | %(
 # Load environment variables from .env file
 load_dotenv()
 
-# Tesseract language options
-langs = os.popen('tesseract --list-langs').read().split('\n')[1:-1]
-
 # Initialize global variables
 vector_index = None
 query_log = []
 sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
 
-def extract_text_from_pdf(pdf_path, lang=None):
-    text = ""
-    image_count = 0
-    total_pages = 0
+langs = os.popen('tesseract --list-langs').read().split('\n')[1:-1]
+
+def preprocess_image(image_path):
+    """
+    Pre-process the image to improve OCR results.
+    - Convert to grayscale
+    - Apply thresholding to improve contrast
+    - Apply denoising if needed
+    """
+    img = cv2.imread(image_path)
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    gray = cv2.equalizeHist(gray)
+    gray = cv2.GaussianBlur(gray, (5, 5), 0)
+    processed_image = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+                                            cv2.THRESH_BINARY, 11, 2)
+    temp_filename = "processed_image.png"
+    cv2.imwrite(temp_filename, processed_image)
+    return temp_filename
+
+def extract_text_from_image(image_path, lang='eng'):
+    processed_image_path = preprocess_image(image_path)
+    text = pytesseract.image_to_string(Image.open(processed_image_path), lang=lang)
+    return text
 
+def extract_text_from_pdf(pdf_path, lang='eng'):
+    text = ""
     try:
         with open(pdf_path, 'rb') as file:
             pdf_reader = PyPDF2.PdfReader(file)
-            total_pages = len(pdf_reader.pages)
-
-            for page_num, page in enumerate(pdf_reader.pages, 1):
+            for page_num in range(len(pdf_reader.pages)):
+                page = pdf_reader.pages[page_num]
                 page_text = page.extract_text()
-
-                # If text is not found, consider the page as an image and use OCR
-                if not page_text.strip():
-                    images = convert_from_path(pdf_path, first_page=page_num, last_page=page_num)
-                    for image in images:
-                        ocr_text = pytesseract.image_to_string(image, lang=None if lang == [] else lang)
-                        text += ocr_text
-                        image_count += 1
-                        text += f"\n[OCR applied on image detected on page {page_num}]\n"
-                else:
+                if page_text.strip():
                     text += page_text
+                else:
+                    images = convert_from_path(pdf_path, first_page=page_num + 1, last_page=page_num + 1)
+                    for image in images:
+                        image.save('temp_image.png', 'PNG')
+                        text += extract_text_from_image('temp_image.png', lang=lang)
+                        text += f"\n[OCR applied on page {page_num + 1}]\n"
     except Exception as e:
-        logging.error(f"Error processing PDF {pdf_path}: {str(e)}")
-        return f"[Error processing PDF: {str(e)}]\n"
-
-    if image_count == total_pages:
-        summary = f"This document consists of {total_pages} page(s) of images.\n"
-        summary += "No text could be extracted directly. OCR was applied to images.\n"
-        summary += f"File path: {pdf_path}\n"
-        return summary
-    elif image_count > 0:
-        text = f"This document contains both text and images.\n" + \
-               f"Total pages: {total_pages}\n" + \
-               f"Pages with images: {image_count}\n" + \
-               f"Extracted text (including OCR):\n\n" + text
-
+        return f"Error processing PDF: {str(e)}"
     return text
 
-def load_docx_file(docx_path):
-    try:
-        doc = docx.Document(docx_path)
-        return '\n'.join([para.text for para in doc.paragraphs])
-    except Exception as e:
-        logging.error(f"Error processing DOCX {docx_path}: {str(e)}")
-        return f"[Error processing DOCX: {str(e)}]\n"
-
-def load_txt_file(txt_path):
-    try:
-        with open(txt_path, 'r', encoding='utf-8') as f:
-            return f.read()
-    except Exception as e:
-        logging.error(f"Error processing TXT {txt_path}: {str(e)}")
-        return f"[Error processing TXT: {str(e)}]\n"
-
-def load_file_based_on_extension(file_path, lang=None):
-    if file_path.lower().endswith('.pdf'):
+def extract_text(file_path, lang='eng'):
+    file_ext = file_path.lower().split('.')[-1]
+    if file_ext in ['pdf']:
         return extract_text_from_pdf(file_path, lang)
-    elif file_path.lower().endswith('.docx'):
-        return load_docx_file(file_path)
-    elif file_path.lower().endswith('.txt'):
-        return load_txt_file(file_path)
+    elif file_ext in ['png', 'jpg', 'jpeg']:
+        return extract_text_from_image(file_path, lang)
     else:
-        raise ValueError(f"Unsupported file format: {file_path}")
+        return f"Unsupported file type: {file_ext}"
 
 def process_upload(api_key, files, lang):
     global vector_index
@@ -111,7 +93,7 @@ def process_upload(api_key, files, lang):
 
     for file_path in files:
         try:
-            text = load_file_based_on_extension(file_path, lang)
+            text = extract_text(file_path, lang)
             if "This document consists of" in text and "page(s) of images" in text:
                 image_heavy_docs.append(os.path.basename(file_path))
             documents.append(Document(text=text))
@@ -137,30 +119,18 @@ def process_upload(api_key, files, lang):
     else:
         return f"No valid documents were indexed. Errors: {'; '.join(error_messages)}", None
 
-# Define the calculate_similarity function
 def calculate_similarity(response, ground_truth):
-    # Encode the response and ground truth
     response_embedding = sentence_model.encode(response, convert_to_tensor=True)
     truth_embedding = sentence_model.encode(ground_truth, convert_to_tensor=True)
 
-    # Convert embeddings to numpy arrays for easier manipulation
-    response_embedding = response_embedding.cpu().numpy()
-    truth_embedding = truth_embedding.cpu().numpy()
-
-    # Normalize the embeddings to unit vectors (magnitude of 1)
     response_embedding = response_embedding / np.linalg.norm(response_embedding)
     truth_embedding = truth_embedding / np.linalg.norm(truth_embedding)
 
-    # Calculate cosine similarity using numpy's dot product
     similarity = np.dot(response_embedding, truth_embedding)
-
-    # Return similarity as a percentage (between 0 and 100)
-    similarity_percentage = (similarity + 1) / 2 * 100  # Normalize from [-1, 1] to [0, 100]
+    similarity_percentage = (similarity + 1) / 2 * 100
 
     return similarity_percentage
 
-
-# This is the missing query_app function that needs to be defined
 def query_app(query, model_name, use_similarity_check, openai_api_key):
     global vector_index, query_log