updated app

app.py

@@ -1,4 +1,24 @@
+import subprocess
+import sys
 import streamlit as st
+
+def install_package(package, flags=False):
+    try:
+        __import__(package)
+        print(f"'{package}' is already installed.")
+    except ImportError:
+        print(f"Installing '{package}'...")
+        if flags:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-build-isolation", package])
+        else:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+if "installed_libraries" not in st.session_state or not st.session_state.installed_libraries:
+    install_package("flash-attn", flags=True)
+    install_package("rerankers[all]")
+    st.session_state.installed_libraries = True
+
+
 from PIL import Image
 import os
 import torch
@@ -11,7 +31,7 @@ from src.pipelines import InvoiceGenerator
 st.set_page_config(page_title="Invoice generator", layout="wide")
 output_folder = "output"
 data_folder = "data"
-template = "template.tex"
+template = "template.md"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 with open(f"{data_folder}/car_parts.json", "r") as f:
     car_parts = json.load(f)
@@ -23,6 +43,7 @@ if "pipeline" not in st.session_state:
         image_invoice_index_path=f"{data_folder}/image_invoice.csv",
         path_to_invoices=f"{data_folder}/invoices",
         path_to_images=f"{data_folder}/images",
+        path_to_template=f"{data_folder}/{template}",
         reranker_model="monovlm",
         device=device,
         gpu_memory_utilization=0.65
@@ -44,24 +65,24 @@ def read_markdown_file(file_path):
 
 
 def display_invoice(image_path):
-    output_md = "invoice_" + os.path.basename(image_path).split(".")[0] + ".md"
-    path_to_output_md = f"{output_folder}/{output_md}"
+    output_pdf = "invoice_" + os.path.basename(image_path).split(".")[0] + ".pdf"
+    path_to_output_pdf = f"{output_folder}/{output_pdf}"
     try:
         result = pipeline.generate_invoice(
-            image_path=image_path, output_path=path_to_output_md, car_parts=car_parts
+            image_path=image_path, output_path=path_to_output_pdf, car_parts=car_parts
         )
         if result is None:
             st.write("Image is irrelevant, upload another one")
             st.session_state["status"] = "irrelevant"
             return
-        print(f"Generated invoice: {path_to_output_md}")
-        st.session_state["generated_md"] = path_to_output_md
+        print(f"Generated invoice: {path_to_output_pdf}")
+        st.session_state["generated_pdf"] = path_to_output_pdf
         st.session_state["invoice_info"] = result["invoice_info"]
         st.session_state["invoice_path"] = f"{data_folder}/invoices/{result['invoice_path']}"
         st.session_state["similar_image"] = f"{data_folder}/images/{result['similar_image']}"
         st.session_state["damage_description"] = result["damage_description"]
         st.session_state["detailed_damage_description"] = result["detailed_damage_description"]
-        return read_markdown_file(path_to_output_md)
+        return get_image_from_pdf(path_to_output_pdf)
     except Exception as e:
         st.write("Could not generate invoice, please try again")
         print(e)
@@ -106,8 +127,7 @@ with col2:
 
 with col3:
     if st.session_state.get("status") == "loaded":
-        # st.image(st.session_state["invoice"], caption="Generated invoice", use_container_width=True)
-        st.markdown(st.session_state["invoice"])
+        st.image(st.session_state["invoice"], caption="Generated invoice", use_container_width=True)
         st.image(st.session_state["similar_image"], caption="Similar accident", width=300)
         st.write(f"Detailed damage description: {st.session_state['detailed_damage_description']}")
         st.write(f"Damage description: {st.session_state['damage_description']}")

src/pipelines.py

@@ -4,6 +4,7 @@ from src.prompts import GENERATE_INVOICE_PROMPT, GENERATE_BRIEF_DAMAGE_DESCRIPTI
     GENERATE_DETAILED_DAMAGE_DESCRIPTION_PROMPT
 from rerankers import Reranker
 import re
+from md2pdf.core import md2pdf
 from fuzzywuzzy import fuzz
 
 
@@ -14,6 +15,7 @@ class InvoiceGenerator:
             image_invoice_index_path,
             path_to_invoices,
             path_to_images,
+            path_to_template,
             reranker_model=None,
             device="cuda",
             max_model_len=4096, max_tokens=2048, gpu_memory_utilization=0.95
@@ -33,15 +35,13 @@ class InvoiceGenerator:
         )
         self.path_to_invoices = path_to_invoices
         self.path_to_images = path_to_images
+        self.path_to_template = path_to_template
 
-    def format_invoice(self, generated_invoice, output_path, template_path="data/template.md"):
-        with open(template_path, "r") as f:
+    def format_invoice(self, generated_invoice, output_path):
+        with open(self.path_to_template, "r") as f:
             md_text = f.read()
         md_text = md_text.replace(r"<<table>>", generated_invoice)
-        with open(output_path, "w") as f:
-            f.write(md_text)
-        #  md2pdf(output_path, md_content=md_text)
-        return md_text
+        md2pdf(output_path, md_content=md_text)
 
     @staticmethod
     def check_within_range(generated_invoice, car_parts):
@@ -145,6 +145,6 @@ class InvoiceGenerator:
             print(comparing_results)
 
         if output_path:
-            self.format_invoice(generated_invoice=generated_invoice, output_path=output_path, template_path=template_path)
+            self.format_invoice(generated_invoice=generated_invoice, output_path=output_path)
 
         return result