passport-processing/ui.py at main · shafed/passport-processing · GitHub

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
import io

import fitz
import pandas as pd
import streamlit as st
from PIL import Image

from make_barcode import generate_barcodes, return_barcode


# Convert pdf to image because streamlit doesn't support pdf view
def pdf_to_images(uploaded_pdf):
    pdf_bytes = uploaded_pdf.read()
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    images = []
    for page in doc:
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
        img_bytes = pix.tobytes("png")
        image = Image.open(io.BytesIO(img_bytes))
        images.append(image)
    return images


def view_and_correct(result_of_OCR, index):
    table = pd.DataFrame(result_of_OCR)
    edited_table = st.data_editor(table, use_container_width=True, key=index)
    edited_table["passport_id"] = f"PAS-{index + 1:06}"
    return edited_table


def convert_for_download(df):
    return df.to_csv().encode("utf-8")


all_tables = []
passport_map = {}
if "passport_map" not in st.session_state:
    st.session_state.passport_map = {}

# Page config
st.set_page_config(page_title="Паспорта оборудования", layout="wide")
tab_digitize, tab_lookup = st.tabs(["Оцифровка", "Поиск по базе"])
with tab_digitize:
    # Upload area
    uploaded_files = st.file_uploader(
        "Загрузите файлы",
        accept_multiple_files=True,
        type=["jpg", "jpeg", "png", "pdf"],
    )
    if uploaded_files:
        for i, uploaded_file in enumerate(uploaded_files):
            # Two columns for scan and info
            col_doc, col_card = st.columns([1, 1])
            with col_doc:
                with st.expander("Развернуть изображение", expanded=False):
                    if uploaded_file.type.startswith("image/"):
                        st.image(uploaded_file, use_container_width=True)
                    elif uploaded_file.type == "application/pdf":
                        images = pdf_to_images(uploaded_file)
                        for page_number, image in enumerate(images, start=1):
                            st.write(f"Страница {page_number}")
                            st.image(image, use_container_width=True)
            with col_card:
                edited_table = view_and_correct(  # TODO использую заглушку
                    [
                        {
                            "serial_number": "12345",
                            "manufacturer": "ООО Завод",
                            "date": "2024-01-01",
                        }
                    ],
                    i,
                )
                st.image(return_barcode(f"PAS-{i + 1:06}"), width=100)
            all_tables.append(edited_table)
            passport_id = f"PAS-{i + 1:06}"
            st.session_state.passport_map[passport_id] = edited_table

        final_table = pd.concat(all_tables, ignore_index=True)

        # Create button
        csv = convert_for_download(final_table)
        st.download_button(
            label="Экспортировать таблицу как CSV",
            data=csv,
            file_name="passport.csv",
            mime="text/csv",
            icon=":material/download:",
        )
        if st.button("Сгенерировать штрихкод", type="primary"):
            generate_barcodes(len(uploaded_files), "PAS", "barcodes/passports")
            st.success("Штрихкоды сгенерированы")

with tab_lookup:
    st.title("Поиск паспорта по штрихкоду")
    code = st.text_input("Впешите значение штрихкода", key="scanner_input")
    if code:
        code = code.strip()
        if code in st.session_state.passport_map:
            passport = st.session_state.passport_map[code]
            st.success(f"Найден паспорт {code}")
            st.dataframe(passport, use_container_width=True)
        else:
            st.error("Паспорт не найден")