| import streamlit as st | |
| from PIL import Image | |
| import os | |
| import torch | |
| from uuid import uuid4 | |
| import fitz | |
| import json | |
| from src.pipelines import InvoiceGenerator | |
| st.set_page_config(page_title="Invoice generator", layout="wide") | |
| output_folder = "output" | |
| data_folder = "data" | |
| template = "template.tex" | |
| 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) | |
| # 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) | |