File size: 6,988 Bytes
3d02141 f11a85d 3d02141 4e6f2f3 3d02141 a578eb1 4e6f2f3 77ecec5 4e6f2f3 6e8796e 2093584 3ec40b4 4e6f2f3 a303cd3 3d02141 a303cd3 0ca755d c6179a2 a303cd3 e8b6d9f 7a67260 a578eb1 7a67260 a303cd3 3d02141 c4c635a 596a523 c4c635a 596a523 c4c635a a303cd3 c6179a2 a303cd3 c6179a2 a303cd3 596a523 a303cd3 596a523 a303cd3 7c35d00 |
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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
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"]))
|