File size: 4,417 Bytes
f11a85d 4e6f2f3 b2136cd |
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 |
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)
|