From 220668c27a6037e8fec60e90cda92e1bfc2e9ad9 Mon Sep 17 00:00:00 2001
From: Amira24523 <85802878+Amira24523@users.noreply.github.com>
Date: Thu, 12 Dec 2024 01:48:19 +0100
Subject: [PATCH] ES FUNKTIONIERT

---
 lib.py            | 17 +++++++++++++++++
 main.py           |  2 +-
 test_PDFandJPG.py |  4 ++--
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/lib.py b/lib.py
index 90d3318..883cd5a 100644
--- a/lib.py
+++ b/lib.py
@@ -5,6 +5,8 @@ import psycopg2
 import pymongo
 import pytesseract
 import ollama
+import io
+from PIL import Image
 from aiohttp.web_response import json_response
 
 import ollama_conn
@@ -47,6 +49,21 @@ def extract_image_content(path):
     img = cv2.imread(path)
     return pytesseract.image_to_string(img)
 
+def extract_text_and_pictures(pdf_datei):
+    pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
+    doc = pymupdf.open(pdf_datei)
+    text = ""
+    for seite in doc:
+        text += seite.get_text()
+        for bild in seite.get_images(full=True):
+            xref = bild[0]
+            basis_bild = doc.extract_image(xref)
+            bild_daten = basis_bild["image"]
+            bild_objekt = Image.open(io.BytesIO(bild_daten))
+            bild_text = pytesseract.image_to_string(bild_objekt, lang="deu")
+            text += " " + bild_text
+    return text
+
 def prompt_embedding(prompt):
     embedding = ollama.embeddings(model="mxbai-embed-large", prompt=prompt)
 
diff --git a/main.py b/main.py
index 8c586a1..1b3b00a 100644
--- a/main.py
+++ b/main.py
@@ -12,7 +12,7 @@ count = 0
 
 def retrieve_file_contents(path):
     file_extraction_functions = {
-        "pdf": lambda path: lib.extract_pdf_content(path),
+        "pdf": lambda path: lib.extract_text_and_pictures(path),
         "jpg": lambda path: lib.extract_image_content(path),
         "png": lambda path: lib.extract_image_content(path),
         "txt": lambda path: lib.extract_pdf_content(path),
diff --git a/test_PDFandJPG.py b/test_PDFandJPG.py
index 580f23e..7073135 100644
--- a/test_PDFandJPG.py
+++ b/test_PDFandJPG.py
@@ -5,7 +5,7 @@ import pytesseract
 
 pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
 
-def extrahiere_text_und_bilder(pdf_datei):
+def extract_text_and_pictures(pdf_datei):
     doc = pymupdf.open(pdf_datei)
     text = ""
     for seite in doc:
@@ -20,6 +20,6 @@ def extrahiere_text_und_bilder(pdf_datei):
     return text
 
 pdf_datei = 'any.pdf'
-text = extrahiere_text_und_bilder(pdf_datei)
+text = extract_text_and_pictures(pdf_datei)
 
 print(text)
\ No newline at end of file