Update app/vector_store.py

app/vector_store.py

@@ -1,10 +1,18 @@
+from PIL import Image
+import pytesseract
+from pdf2image import convert_from_path
 import os
 from langchain.vectorstores import Chroma
 from langchain.embeddings import SentenceTransformerEmbeddings
 from langchain.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.docstore.document import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
 
 CHROMA_DIR = "data/chroma_db"
+CHROMA_IMG_DIR = "data/image_db"
 
 
 def store_pdf(pdf_path):
@@ -15,3 +23,24 @@ def store_pdf(pdf_path):
     chunks = splitter.split_documents(docs)
     embeddings = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
     Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_DIR)
+
+
+def store_pdf_image(pdf_path):
+    text = extract_text_from_scanned_pdf(pdf_path)
+    doc = Document(page_content=text)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    chunks = splitter.split_documents([doc])
+    embeddings = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
+    Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_IMG_DIR)
+
+
+# images = convert_from_path("your_file.pdf", poppler_path="/opt/homebrew/bin")
+
+
+def extract_text_from_scanned_pdf(pdf_path):
+    pages = convert_from_path(pdf_path, dpi=300)
+    all_text = ""
+    for i, page in enumerate(pages):
+        text = pytesseract.image_to_string(page, lang="eng")
+        all_text += f"\n--- Page {i+1} ---\n{text}"
+    return all_text