forked from microsoft/table-transformer
-
Notifications
You must be signed in to change notification settings - Fork 7
/
app.py
51 lines (32 loc) · 1.08 KB
/
app.py
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
from datetime import datetime
import io
import cv2
import numpy as np
import streamlit as st
from PIL import Image
st.set_page_config(layout="wide")
st.sidebar.header("TABLE RECOGNITION DEMO")
uploaded_file = st.file_uploader(label="UPLOAD FILE", type=["png", "jpg", "jpeg"])
image = None
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
from core import TableRecognizer
m = TableRecognizer(
checkpoint_path="/data/pubtables1m/PubTables1M-Structure-PASCAL-VOC/output/20220110130310/model_11.pth"
)
def main():
global image
if image is None:
return
placeholder = st.image(image, width=500) # display image
with st.spinner("🤖 AI is at Work! "):
start_time = datetime.now()
results = m.predict(image_path=image)
output_image = results["debug_image"]
consume_time = datetime.now() - start_time
st.text(f"Consume time: {consume_time}")
placeholder.empty()
st.image([image, output_image], caption=["input", "output"], width=500)
st.balloons()
if __name__ == "__main__":
main()