updated app

app.py

@@ -18,93 +18,93 @@ with open(f"{data_folder}/car_parts.json", "r") as f:
     car_parts = json.load(f)
 
 
-# if "pipeline" not in st.session_state:
-#     st.session_state.pipeline = InvoiceGenerator(
-#         fais_index_path=f"{data_folder}/invoice_index.faiss",
-#         image_invoice_index_path=f"{data_folder}/image_invoice.csv",
-#         path_to_invoices=f"{data_folder}/invoices",
-#         path_to_images=f"{data_folder}/images",
-#         reranker_model="monovlm",
-#         device=device,
-#         gpu_memory_utilization=0.65
-#     )
-# pipeline = st.session_state.pipeline
-#
-#
-# def get_image_from_pdf(pdf_path):
-#     doc = fitz.open(pdf_path)
-#     page = doc[0]
-#     mat = fitz.Matrix(2, 2)
-#     pix = page.get_pixmap(matrix=mat)
-#     img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-#     return img
-#
-#
-# 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:
-#         result = pipeline.generate_invoice(
-#             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_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 get_image_from_pdf(path_to_output_pdf)
-#     except Exception as e:
-#         st.write("Could not generate invoice, please try again")
-#         print(e)
-#         return None
-#
-#
-# st.title("Upload FNOL photo")
-#
-# col1, col2, col3 = st.columns([4, 1, 4])
-#
-# with col1:
-#     uploaded_image = st.file_uploader("Upload photo", type=["jpg", "jpeg", "png"])
-#     if uploaded_image:
-#         try:
-#             image = Image.open(uploaded_image)
-#             image_path = f"{output_folder}/{str(uuid4())[:5]}.png"
-#             image.save(image_path)
-#             print(f"Image: {image_path}")
-#             st.image(image, caption="Uploaded photo", width=300)
-#             st.session_state["image"] = image_path
-#         except Exception as e:
-#             st.write(f"Coudn't load image: {e}")
-#
-# with col2:
-#     if st.session_state.get("image"):
-#         if st.button("Generate invoice"):
-#             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)
-#     if st.session_state.get("generated_pdf"):
-#         with open(st.session_state["generated_pdf"], "rb") as f:
-#             file_data = f.read()
-#         st.download_button(
-#             label="Download invoice",
-#             data=file_data,
-#             file_name="generated_invoice.pdf",
-#             mime="application/pdf"
-#         )
-#
-# with col3:
-#     if st.session_state.get("status") == "loaded":
-#         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']}")
-#         st.image(get_image_from_pdf(st.session_state["invoice_path"]), caption="Invoice of similar accident type", use_container_width=True)
+if "pipeline" not in st.session_state:
+    st.session_state.pipeline = InvoiceGenerator(
+        fais_index_path=f"{data_folder}/invoice_index.faiss",
+        image_invoice_index_path=f"{data_folder}/image_invoice.csv",
+        path_to_invoices=f"{data_folder}/invoices",
+        path_to_images=f"{data_folder}/images",
+        reranker_model="monovlm",
+        device=device,
+        gpu_memory_utilization=0.65
+    )
+pipeline = st.session_state.pipeline
+
+
+def get_image_from_pdf(pdf_path):
+    doc = fitz.open(pdf_path)
+    page = doc[0]
+    mat = fitz.Matrix(2, 2)
+    pix = page.get_pixmap(matrix=mat)
+    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    return img
+
+
+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:
+        result = pipeline.generate_invoice(
+            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_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 get_image_from_pdf(path_to_output_pdf)
+    except Exception as e:
+        st.write("Could not generate invoice, please try again")
+        print(e)
+        return None
+
+
+st.title("Upload FNOL photo")
+
+col1, col2, col3 = st.columns([4, 1, 4])
+
+with col1:
+    uploaded_image = st.file_uploader("Upload photo", type=["jpg", "jpeg", "png"])
+    if uploaded_image:
+        try:
+            image = Image.open(uploaded_image)
+            image_path = f"{output_folder}/{str(uuid4())[:5]}.png"
+            image.save(image_path)
+            print(f"Image: {image_path}")
+            st.image(image, caption="Uploaded photo", width=300)
+            st.session_state["image"] = image_path
+        except Exception as e:
+            st.write(f"Coudn't load image: {e}")
+
+with col2:
+    if st.session_state.get("image"):
+        if st.button("Generate invoice"):
+            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)
+    if st.session_state.get("generated_pdf"):
+        with open(st.session_state["generated_pdf"], "rb") as f:
+            file_data = f.read()
+        st.download_button(
+            label="Download invoice",
+            data=file_data,
+            file_name="generated_invoice.pdf",
+            mime="application/pdf"
+        )
+
+with col3:
+    if st.session_state.get("status") == "loaded":
+        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']}")
+        st.image(get_image_from_pdf(st.session_state["invoice_path"]), caption="Invoice of similar accident type", use_container_width=True)