from datasets import load_dataset
import streamlit as st
from mmocr.apis import MMOCRInferencer
import mmcv
import matplotlib.pyplot as plt

st.write('Hello, *World!* :sunglasses:')
dataset = load_dataset("katanaml-org/invoices-donut-data-v1")
st.write(f"{len(dataset)}")

infer = MMOCRInferencer(det='dbnetpp', rec='svtr-small', kie='SDMGR')
result = infer('demo/demo_kie.jpeg', save_vis=True)

predicted_img = mmcv.imread('results/vis/demo_kie.jpg')
plt.figure(figsize=(18, 32))
plt.imshow(mmcv.bgr2rgb(predicted_img))
plt.show()