added granularity

app.py

@@ -83,8 +83,12 @@ def display_invoice(image_path):
     output_pdf = "invoice_" + os.path.basename(image_path).split(".")[0] + ".pdf"
     path_to_output_pdf = f"{output_folder}/{output_pdf}"
     try:
+        print(f"Generate granular invoice: {st.session_state['granular_invoice']}")
         result = pipeline.generate_invoice(
-            image_path=image_path, output_path=path_to_output_pdf, car_parts=car_parts
+            image_path=image_path,
+            output_path=path_to_output_pdf,
+            car_parts=car_parts,
+            use_granular_invoice=st.session_state["granular_invoice"]
         )
         if result is None:
             st.write("Image is irrelevant, upload another one")
@@ -124,13 +128,21 @@ with col1:
 
 with col2:
     if st.session_state.get("image"):
-        if st.button("Generate invoice"):
+        if st.button("Generate regular invoice"):
+            st.session_state["granular_invoice"] = False
+            with st.spinner("Generating..."):
+                st.session_state["invoice"] = display_invoice(st.session_state["image"])
+                if st.session_state["invoice"]:
+                    st.session_state["status"] = "loaded"
+        if st.button("Generate granular invoice"):
+            st.session_state["granular_invoice"] = True
             with st.spinner("Generating..."):
                 st.session_state["invoice"] = display_invoice(st.session_state["image"])
                 if st.session_state["invoice"]:
                     st.session_state["status"] = "loaded"
     else:
-        st.button("Generate invoice", disabled=True)
+        st.button("Generate regular invoice", disabled=True)
+        st.button("Generate granular invoice", disabled=True)
     if st.session_state.get("generated_pdf"):
         with open(st.session_state["generated_pdf"], "rb") as f:
             file_data = f.read()

src/RAG.py

@@ -34,8 +34,17 @@ class RAG:
         self.path_to_images = path_to_images
         self.reranker = reranker
         if path_to_invoice_json:
-            with open(path_to_invoice_json, "r") as f:
-                self.invoice_json = json.load(f)
+            if type(path_to_invoice_json) == str:
+                with open(path_to_invoice_json, "r") as f:
+                    self.invoice_json = json.load(f)
+            elif type(path_to_invoice_json) == dict and set(list(path_to_invoice_json.keys())) == {"invoices", "invoices_granular"}:
+                with open(path_to_invoice_json["invoices"], "r") as f:
+                    self.invoice_json = json.load(f)
+                with open(path_to_invoice_json["invoices_granular"], "r") as f:
+                    self.invoice_json_granular = json.load(f)
+            else:
+                raise ValueError("Invalid format for invoice json.")
+
 
     @staticmethod
     def image_to_base64(image_path):
@@ -61,14 +70,18 @@ class RAG:
         distances, indices = self.index.search(image_features, k)
         return distances, indices
 
-    def return_invoice_table(self, path=None, invoice_is_table=True):
+    def return_invoice_table(self, path=None, invoice_is_table=True, use_granular_invoice=False):
         if path is None and not invoice_is_table:
             raise ValueError("Path to invoice must be provided.")
         if self.invoice_json is None and invoice_is_table:
             raise ValueError("Path to invoice json must be provided.")
+        if self.invoice_json_granular is None and use_granular_invoice:
+            raise ValueError("Path to granular invoice json must be provided.")
 
-        if invoice_is_table:
+        if invoice_is_table and not use_granular_invoice:
             return self.invoice_json[path]
+        elif invoice_is_table and use_granular_invoice:
+            return self.invoice_json_granular[path]
 
         pdf_path = f"{self.path_to_invoices}/{path}"
         reader = PdfReader(pdf_path)
@@ -98,7 +111,8 @@ class RAG:
             return_only_path=True,
             k=1,
             damage_description=None,
-            invoice_is_table=True
+            invoice_is_table=True,
+            use_granular_invoice=False
     ):
         if self.image_invoice_index is None:
             raise ValueError("No index for invoices found.")
@@ -124,7 +138,9 @@ class RAG:
         invoices_tables = []
 
         for invoice in invoices:
-            reformatted_inv = self.return_invoice_table(invoice, invoice_is_table)
+            reformatted_inv = self.return_invoice_table(
+                invoice, invoice_is_table, use_granular_invoice=use_granular_invoice
+            )
             invoices_tables.append(reformatted_inv)
 
         return invoices_tables, invoices

src/pipelines.py

@@ -106,7 +106,9 @@ class InvoiceGenerator:
         all_lines[first_cost_line:last_cost_line] = list(map(lambda x: " | ".join(x), cost_lines))
         return "\n".join(all_lines)
 
-    def generate_invoice(self, image_path, output_path=None, car_parts=None):
+    def generate_invoice(
+            self, image_path, output_path=None, car_parts=None, use_granular_invoice=False
+    ):
 
         result = {}
 
@@ -121,7 +123,11 @@ class InvoiceGenerator:
         print("Damage Description:", damage_description)
 
         invoice_info, invoice_path = self.rag.find_invoice(
-            image_path=image_path, return_only_path=False, damage_description=damage_description, k=5
+            image_path=image_path,
+            return_only_path=False,
+            damage_description=damage_description,
+            k=5,
+            use_granular_invoice=use_granular_invoice
         )
         invoice_info = invoice_info[0]
         invoice_path = invoice_path[0]