Update services/extract_text.py

services/extract_text.py

@@ -1,3 +1,168 @@
+# import os
+# import logging
+# import fitz  # PyMuPDF
+# import numpy as np
+# from PIL import Image
+# import cv2
+# import re
+
+# # OCR
+# from paddleocr import PaddleOCR
+
+# # Optional Mistral OCR
+# try:
+#     from doctr.models import ocr_predictor
+#     from doctr.io import DocumentFile
+#     mistral_ocr = ocr_predictor(pretrained=True)
+#     use_mistral_ocr = True
+# except ImportError:
+#     mistral_ocr = None
+#     use_mistral_ocr = False
+
+# # Environment paths
+# os.environ.setdefault("HOME", "/app")
+# os.environ.setdefault("PADDLEOCR_HOME", "/app/.paddleocr")
+
+# # Logging
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger(__name__)
+
+# # PaddleOCR
+# ocr = PaddleOCR(use_angle_cls=True, lang='en')
+
+# def clean_text(text):
+#     return re.sub(r'\s+', ' ', text).strip()
+
+# def auto_rotate_image(pil_img):
+#     """Auto-rotate PIL image safely."""
+#     if pil_img.mode != "RGB":
+#         pil_img = pil_img.convert("RGB")
+#     img_cv = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2GRAY)
+#     coords = np.column_stack(np.where(img_cv > 0))
+#     if coords.size == 0:
+#         return pil_img  # blank page
+#     angle = cv2.minAreaRect(coords)[-1]
+#     angle = -(90 + angle) if angle < -45 else -angle
+#     (h, w) = img_cv.shape[:2]
+#     M = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
+#     rotated = cv2.warpAffine(img_cv, M, (w, h),
+#                              flags=cv2.INTER_CUBIC,
+#                              borderMode=cv2.BORDER_REPLICATE)
+#     return Image.fromarray(cv2.cvtColor(rotated, cv2.COLOR_GRAY2RGB))
+
+# def extract_images_with_fitz(pdf_path, start_page=1, end_page=None):
+#     images = []
+#     try:
+#         doc = fitz.open(pdf_path)
+#         total_pages = len(doc)
+#         end = min(end_page or total_pages, total_pages)
+#         for i in range(start_page - 1, end):
+#             try:
+#                 page = doc[i]
+#                 pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
+#                 mode = "RGBA" if pix.alpha else "RGB"
+#                 img = Image.frombytes(mode, [pix.width, pix.height], pix.samples)
+#                 images.append((i + 1, img))
+#             except Exception as e:
+#                 logger.error(f"Error rendering page {i + 1}: {e}")
+#         doc.close()
+#     except Exception as e:
+#         logger.error(f"Failed to open PDF file: {e}")
+#     return images
+
+# def extract_text_from_file(file, start_page=None, end_page=None, filename=None):
+#     ext = os.path.splitext(filename or "")[-1].lower()
+#     result = []
+
+#     if ext == ".pdf":
+#         try:
+#             doc = fitz.open(file.name)
+#         except Exception as e:
+#             logger.error(f"Cannot open PDF {filename}: {e}")
+#             return "[Error opening PDF]"
+
+#         images = extract_images_with_fitz(file.name, start_page or 1, end_page)
+#         total_pages = len(doc)
+#         start = max(start_page or 1, 1)
+#         end = min(end_page or total_pages, total_pages)
+
+#         for i, page in enumerate(doc):
+#             page_num = i + 1
+#             if not (start <= page_num <= end):
+#                 continue
+
+#             text = page.get_text()
+#             if text.strip():
+#                 result.append(f"Page {page_num} (Extracted):\n{clean_text(text)}")
+#             else:
+#                 if i < len(images):
+#                     try:
+#                         img = auto_rotate_image(images[i][1])
+#                         img_np = np.array(img)
+#                         ocr_text = ""
+#                         # PaddleOCR
+#                         try:
+#                             ocr_result = ocr.ocr(img_np, cls=True)
+#                             ocr_text = "\n".join([line[1][0] for line in ocr_result[0]]) if ocr_result else ""
+#                         except Exception as e:
+#                             logger.warning(f"PaddleOCR failed on page {page_num}: {e}")
+
+#                         # Mistral OCR fallback
+#                         if not ocr_text and use_mistral_ocr:
+#                             try:
+#                                 doc_img = DocumentFile.from_images(img)
+#                                 ocr_text = mistral_ocr(doc_img).render()
+#                             except Exception as e:
+#                                 logger.warning(f"Mistral OCR failed on page {page_num}: {e}")
+#                                 ocr_text = "[OCR Error]"
+
+#                         result.append(f"Page {page_num} (OCR):\n{clean_text(ocr_text) or '[No OCR Text]'}")
+#                     except Exception as e:
+#                         logger.error(f"OCR processing failed for page {page_num}: {e}")
+#                         result.append(f"Page {page_num}: [OCR Error]")
+#                 else:
+#                     result.append(f"Page {page_num}: [No text or image]")
+
+#         doc.close()
+#         return "\n\n".join(result)
+
+#     elif ext == ".docx":
+#         from docx.api import Document
+#         doc = Document(file.name)
+#         paras = [p.text for p in doc.paragraphs if p.text.strip()]
+#         page_texts = []
+#         page_size = 500
+#         for i in range(0, len(paras), page_size):
+#             page_texts.append("\n".join(paras[i:i + page_size]))
+#         selected_pages = page_texts
+#         if start_page and end_page:
+#             selected_pages = page_texts[start_page - 1:end_page]
+#         return clean_text("\n\n".join(selected_pages))
+
+#     elif ext == ".csv":
+#         import pandas as pd
+#         try:
+#             return pd.read_csv(file.name).to_string(index=False)
+#         except Exception as e:
+#             logger.error(f"CSV read error: {e}")
+#             return "[CSV Read Error]"
+
+#     elif ext in [".xls", ".xlsx"]:
+#         import pandas as pd
+#         try:
+#             xl = pd.ExcelFile(file.name)
+#             return "\n\n".join([
+#                 f"Sheet: {s}\n{xl.parse(s).to_string(index=False)}"
+#                 for s in xl.sheet_names
+#             ])
+#         except Exception as e:
+#             logger.error(f"Excel read error: {e}")
+#             return "[Excel Read Error]"
+
+#     else:
+#         return "[Unsupported file type]"
+
+
 import os
 import logging
 import fitz  # PyMuPDF
@@ -5,6 +170,7 @@ import numpy as np
 from PIL import Image
 import cv2
 import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # OCR
 from paddleocr import PaddleOCR
@@ -70,9 +236,54 @@ def extract_images_with_fitz(pdf_path, start_page=1, end_page=None):
         logger.error(f"Failed to open PDF file: {e}")
     return images
 
+
+# -------------------- Parallel Extraction Wrapper --------------------
+
+def try_pymupdf_text(doc, start, end):
+    """Try extracting text using PyMuPDF native text extraction"""
+    result = []
+    for i in range(start-1, end):
+        page = doc[i]
+        text = page.get_text()
+        if text.strip():
+            result.append(f"Page {i+1}:\n{clean_text(text)}")
+    return "\n\n".join(result)
+
+
+def try_paddleocr(images):
+    """Try OCR using PaddleOCR"""
+    result = []
+    for page_num, img in images:
+        img = auto_rotate_image(img)
+        img_np = np.array(img)
+        try:
+            ocr_result = ocr.ocr(img_np, cls=True)
+            ocr_text = "\n".join([line[1][0] for line in ocr_result[0]]) if ocr_result else ""
+            result.append(f"Page {page_num}:\n{clean_text(ocr_text)}")
+        except Exception as e:
+            logger.warning(f"PaddleOCR failed on page {page_num}: {e}")
+    return "\n\n".join(result)
+
+
+def try_mistralocr(images):
+    """Try OCR using Mistral/Doctr OCR"""
+    if not use_mistral_ocr:
+        return ""
+    result = []
+    for page_num, img in images:
+        try:
+            doc_img = DocumentFile.from_images(img)
+            ocr_text = mistral_ocr(doc_img).render()
+            result.append(f"Page {page_num}:\n{clean_text(ocr_text)}")
+        except Exception as e:
+            logger.warning(f"Mistral OCR failed on page {page_num}: {e}")
+    return "\n\n".join(result)
+
+
+# -------------------- Main Extractor --------------------
+
 def extract_text_from_file(file, start_page=None, end_page=None, filename=None):
     ext = os.path.splitext(filename or "")[-1].lower()
-    result = []
 
     if ext == ".pdf":
         try:
@@ -81,63 +292,40 @@ def extract_text_from_file(file, start_page=None, end_page=None, filename=None):
             logger.error(f"Cannot open PDF {filename}: {e}")
             return "[Error opening PDF]"
 
-        images = extract_images_with_fitz(file.name, start_page or 1, end_page)
         total_pages = len(doc)
         start = max(start_page or 1, 1)
         end = min(end_page or total_pages, total_pages)
+        images = extract_images_with_fitz(file.name, start, end)
 
-        for i, page in enumerate(doc):
-            page_num = i + 1
-            if not (start <= page_num <= end):
-                continue
-
-            text = page.get_text()
-            if text.strip():
-                result.append(f"Page {page_num} (Extracted):\n{clean_text(text)}")
-            else:
-                if i < len(images):
-                    try:
-                        img = auto_rotate_image(images[i][1])
-                        img_np = np.array(img)
-                        ocr_text = ""
-                        # PaddleOCR
-                        try:
-                            ocr_result = ocr.ocr(img_np, cls=True)
-                            ocr_text = "\n".join([line[1][0] for line in ocr_result[0]]) if ocr_result else ""
-                        except Exception as e:
-                            logger.warning(f"PaddleOCR failed on page {page_num}: {e}")
-
-                        # Mistral OCR fallback
-                        if not ocr_text and use_mistral_ocr:
-                            try:
-                                doc_img = DocumentFile.from_images(img)
-                                ocr_text = mistral_ocr(doc_img).render()
-                            except Exception as e:
-                                logger.warning(f"Mistral OCR failed on page {page_num}: {e}")
-                                ocr_text = "[OCR Error]"
-
-                        result.append(f"Page {page_num} (OCR):\n{clean_text(ocr_text) or '[No OCR Text]'}")
-                    except Exception as e:
-                        logger.error(f"OCR processing failed for page {page_num}: {e}")
-                        result.append(f"Page {page_num}: [OCR Error]")
-                else:
-                    result.append(f"Page {page_num}: [No text or image]")
+        tasks = {}
+        with ThreadPoolExecutor() as executor:
+            tasks[executor.submit(try_pymupdf_text, doc, start, end)] = "PyMuPDF"
+            tasks[executor.submit(try_paddleocr, images)] = "PaddleOCR"
+            if use_mistral_ocr:
+                tasks[executor.submit(try_mistralocr, images)] = "MistralOCR"
+
+            results = {}
+            for future in as_completed(tasks):
+                method = tasks[future]
+                try:
+                    text = future.result()
+                    results[method] = text
+                except Exception as e:
+                    logger.error(f"{method} failed: {e}")
+                    results[method] = ""
 
         doc.close()
-        return "\n\n".join(result)
+
+        # Pick the longest text among the methods
+        best_method, best_text = max(results.items(), key=lambda kv: len(kv[1].strip()))
+        logger.info(f"Best extraction chosen: {best_method} (length {len(best_text)})")
+        return best_text or "[No text extracted]"
 
     elif ext == ".docx":
-        from docx.api import Document
+        from docx import Document
         doc = Document(file.name)
         paras = [p.text for p in doc.paragraphs if p.text.strip()]
-        page_texts = []
-        page_size = 500
-        for i in range(0, len(paras), page_size):
-            page_texts.append("\n".join(paras[i:i + page_size]))
-        selected_pages = page_texts
-        if start_page and end_page:
-            selected_pages = page_texts[start_page - 1:end_page]
-        return clean_text("\n\n".join(selected_pages))
+        return clean_text("\n".join(paras))
 
     elif ext == ".csv":
         import pandas as pd