File size: 5,542 Bytes
3d02141 f11a85d 3d02141 4e6f2f3 3d02141 0ca755d 4e6f2f3 77ecec5 4e6f2f3 6e8796e 2093584 3ec40b4 4e6f2f3 a303cd3 3d02141 a303cd3 0ca755d a303cd3 e8b6d9f 7a67260 a303cd3 3d02141 c4c635a a303cd3 3d02141 a303cd3 c4c635a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
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 = f"{data_folder}/invoices.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
)
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:
with open(invoice_json_path, "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:
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"] = 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
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.markdown(get_similar_invoice(st.session_state["invoice_path"]))
|