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 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.md" invoice_json_path = { "invoices": f"{data_folder}/invoices.json", "invoices_granular": f"{data_folder}/invoices_granular.json" } 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) print("Successfully loaded car parts") faiss_index_path = f"{data_folder}/invoice_index.faiss" print("Index file size:", os.path.getsize(faiss_index_path)) 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", path_to_template=f"{data_folder}/{template}", reranker_model="monovlm", device=device, invoice_json_path=invoice_json_path, gpu_memory_utilization=0.65, max_tokens=2048, ) 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 read_markdown_file(file_path): with open(file_path, "r") as f: return f.read() def get_similar_invoice(invoice_name): try: invoice_json = invoice_json_path["invoices_granular"] if st.session_state["granular_invoice"] else invoice_json_path["invoices"] with open(invoice_json, "r") as f: invoice_json = json.load(f) return invoice_json[invoice_name] except Exception as e: print(e) return "" 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, use_granular_invoice=st.session_state["granular_invoice"] ) 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"] = result['invoice_path'] print(f"Invoice path: {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 def resize_image_and_save(img: Image, image_path: str, max_width: int = 800, max_height: int = 800): img.thumbnail((max_width, max_height)) img.save(image_path, "PNG", quality=85) 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" resize_image_and_save(image, 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 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 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() 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": if st.session_state.get("invoice"): st.image(st.session_state["invoice"], caption="Generated invoice", use_container_width=True) else: st.write("Invoice could not be displayed. Please regenerate.") if st.session_state.get("similar_image"): st.image(st.session_state["similar_image"], caption="Similar accident", width=300) else: st.write("No similar image available.") st.write(f"Detailed damage description: {st.session_state.get('detailed_damage_description', 'N/A')}") st.write(f"Damage description: {st.session_state.get('damage_description', 'N/A')}") if st.session_state.get("invoice_path"): st.markdown(get_similar_invoice(st.session_state["invoice_path"]))