File size: 5,360 Bytes
3d02141 f11a85d 3d02141 4e6f2f3 3d02141 0ca755d 4e6f2f3 6e8796e 3ec40b4 4e6f2f3 a303cd3 3d02141 a303cd3 0ca755d a303cd3 e8b6d9f a303cd3 3d02141 a303cd3 3d02141 a303cd3 3d02141 a303cd3 3d02141 a303cd3 3d02141 a303cd3 |
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 |
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)
faiss_index_path = f"{data_folder}/invoice_index.faiss"
print("File size:", os.path.getsize(faiss_index_path))
st.image("data/images/invoice_car-body-scratches-789.png", width=300)
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 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)
|