|
import clip |
|
import faiss |
|
from PIL import Image |
|
from pypdf import PdfReader |
|
import pandas as pd |
|
import re |
|
import os |
|
import fitz |
|
import torch |
|
import numpy as np |
|
from tqdm import tqdm |
|
import base64 |
|
|
|
|
|
class RAG: |
|
def __init__( |
|
self, |
|
fais_index_path, |
|
clip_model="ViT-B/32", |
|
reranker=None, |
|
device="cpu", |
|
image_invoice_index_path=None, |
|
path_to_invoices=None, |
|
path_to_images=None |
|
): |
|
self.index = faiss.read_index(fais_index_path) |
|
self.model, self.preprocess = clip.load(clip_model, device=device) |
|
self.device = device |
|
if image_invoice_index_path: |
|
self.image_invoice_index = pd.read_csv(image_invoice_index_path) |
|
self.path_to_invoices = path_to_invoices |
|
self.path_to_images = path_to_images |
|
self.reranker = reranker |
|
|
|
@staticmethod |
|
def image_to_base64(image_path): |
|
with open(image_path, "rb") as image_file: |
|
return base64.b64encode(image_file.read()) |
|
|
|
def search_text(self, text, k=1): |
|
text_features = self.model.encode_text(clip.tokenize([text]).to(self.device)) |
|
text_features /= text_features.norm(dim=-1, keepdim=True) |
|
text_features = text_features.detach().numpy() |
|
distances, indices = self.index.search(text_features, k) |
|
return distances, indices |
|
|
|
def search_image(self, image=None, image_path=None, k=1): |
|
if image is None and image_path is None: |
|
raise ValueError("Either image or image_path must be provided.") |
|
if image is None: |
|
image = Image.open(image_path) |
|
image_input = self.preprocess(image).unsqueeze(0).to(self.device) |
|
image_features = self.model.encode_image(image_input) |
|
image_features /= image_features.norm(dim=-1, keepdim=True) |
|
image_features = image_features.detach().numpy() |
|
distances, indices = self.index.search(image_features, k) |
|
return distances, indices |
|
|
|
def find_invoice(self, image=None, image_path=None, return_only_path=True, k=1, damage_description=None): |
|
if self.image_invoice_index is None: |
|
raise ValueError("No index for invoices found.") |
|
_, indices = self.search_image(image=image, image_path=image_path, k=k) |
|
img_ids = self.image_invoice_index.iloc[indices[0]]['img_id'].values |
|
invoices = self.image_invoice_index[self.image_invoice_index['img_id'].isin(img_ids)]['invoice'].values.tolist() |
|
images_paths = self.image_invoice_index[self.image_invoice_index['img_id'].isin(img_ids)]['image'].values.tolist() |
|
|
|
if self.reranker: |
|
if damage_description is None: |
|
raise ValueError("Damage description must be provided.") |
|
|
|
|
|
images = [f"{self.path_to_images}/{img_path}" for img_path in images_paths] |
|
results = self.reranker.rank(damage_description, images, doc_ids=invoices) |
|
invoices = [doc.doc_id for doc in results] |
|
print(invoices) |
|
|
|
if return_only_path: |
|
return invoices, images_paths |
|
|
|
if not self.path_to_invoices: |
|
raise ValueError("Path to data must be provided.") |
|
|
|
invoices_tables = [] |
|
|
|
for invoice in invoices: |
|
pdf_path = f"{self.path_to_invoices}/{invoice}" |
|
reader = PdfReader(pdf_path) |
|
page = reader.pages[0] |
|
text = page.extract_text() |
|
|
|
table_text = re.search(r"Beschädigtes Teil.*?Gesamtsumme:.*?EUR", text, re.DOTALL).group() |
|
|
|
lines = table_text.splitlines() |
|
header = lines[0] |
|
other_text = "\n".join(lines[1:]) |
|
cleaned_text = re.sub(r"(?<!\d)\n", " ", other_text) |
|
|
|
table = header + "\n" + cleaned_text |
|
|
|
inv = table.split("\n") |
|
reformatted_inv = "Beschädigtes Teil | Teilkosten (EUR) | Arbeitsstunden | Arbeitskosten (EUR/Stunde) | Gesamtkosten (EUR)\n" + "\n".join( |
|
" ".join(inv[i].split(" ")[:-4]) + " | " + ' | '.join(inv[i].split(" ")[-4:]) for i in |
|
range(1, len(inv) - 1)) + "\n" + inv[-1] |
|
|
|
invoices_tables.append(reformatted_inv) |
|
|
|
return invoices_tables, invoices |
|
|
|
|
|
def build_rag(directory): |
|
invoices = os.listdir(f"{directory}/invoices_validated") |
|
invoices = [i for i in invoices if i.endswith(".pdf")] |
|
|
|
image_invoice = [] |
|
os.makedirs(f"{directory}/images", exist_ok=True) |
|
os.makedirs(f"{directory}/invoices", exist_ok=True) |
|
|
|
for invoice in invoices: |
|
doc = fitz.open(f"{directory}/invoices_validated/{invoice}") |
|
|
|
page = doc[1] |
|
image_list = page.get_images(full=True) |
|
text = page.get_text() |
|
|
|
xref = image_list[0][0] |
|
base_image = doc.extract_image(xref) |
|
image_bytes = base_image["image"] |
|
image_name = invoice.replace(".pdf", ".png") |
|
with open(f"{directory}/images/{image_name}", "wb") as img_file: |
|
img_file.write(image_bytes) |
|
|
|
doc.delete_pages(range(1, doc.page_count)) |
|
doc.save(f"{directory}/invoices/{invoice}") |
|
doc.close() |
|
|
|
image_invoice.append({ |
|
"invoice": invoice, |
|
"image": image_name, |
|
"description": text |
|
}) |
|
|
|
image_invoice = pd.DataFrame(image_invoice) |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
model, preprocess = clip.load("ViT-B/32", device=device) |
|
images = image_invoice["image"].tolist() |
|
|
|
embeddings = [] |
|
image_indices = [] |
|
img_ids = [] |
|
|
|
for idx, img_path in enumerate(tqdm(images)): |
|
image = Image.open(f"{directory}/images/{img_path}") |
|
img_ids.append(idx) |
|
inputs = preprocess(image).unsqueeze(0).to(device) |
|
|
|
with torch.no_grad(): |
|
image_embedding = model.encode_image(inputs) |
|
|
|
image_embedding = image_embedding / image_embedding.norm(dim=-1, keepdim=True) |
|
embeddings.append(image_embedding.cpu().numpy().astype("float32")) |
|
image_indices.append(img_path) |
|
|
|
image_invoice["img_id"] = img_ids |
|
image_invoice.to_csv(f"{directory}/image_invoice.csv", index=False) |
|
|
|
embeddings_np = np.vstack(embeddings) |
|
|
|
dimension = embeddings_np.shape[1] |
|
index = faiss.IndexFlatIP(dimension) |
|
index.add(embeddings_np) |
|
|
|
faiss.write_index(index, f"{directory}/invoice_index.faiss") |
|
|
