img_ocr / app.py
tianchiguaixia's picture
Duplicate from tianchiguaixia/ocr
fbaa7e3
# -*- coding: utf-8 -*-
# time: 2022/10/17 11:22
# file: AI-医学图片OCR.py
import streamlit as st
from ocr.ocr import detect, recognize
from ocr.utils import bytes_to_numpy
import pandas as pd
import os
import cv2
from paddleocr import PPStructure, draw_structure_result, save_structure_res
st.title("AI-医学图片OCR")
def convert_df(df):
# IMPORTANT: Cache the conversion to prevent computation on every rerun
return df.to_csv().encode("gbk")
# 上传图片
uploaded_file = st.sidebar.file_uploader(
'请选择一张图片', type=['png', 'jpg', 'jpeg'])
print('uploaded_file:', uploaded_file)
table_engine = PPStructure(show_log=True)
if uploaded_file is not None:
# To read file as bytes:
# content = cv2.imread(uploaded_file)
# st.write(content)
bytes_data = uploaded_file.getvalue()
# 转换格式
img = bytes_to_numpy(bytes_data, channels='RGB')
option_task = st.sidebar.radio('请选择要执行的任务', ('查看原图', '文本检测'))
if option_task == '查看原图':
st.image(img, caption='原图')
elif option_task == '文本检测':
im_show = detect(img)
st.image(im_show, caption='文本检测后的图片')
base_path = "streamlit_data"
path = os.path.exists(base_path + "/" + uploaded_file.name.split('.')[0])
if st.button('✨ 启动!'):
local_path = base_path + "/" + uploaded_file.name.split('.')[0]
result = table_engine(img)
save_structure_res(result, base_path, uploaded_file.name.split('.')[0])
with st.container():
with st.expander(label="json结果展示", expanded=False):
st.write(result)
for i in os.listdir(local_path):
if ".xlsx" in i:
df = pd.read_excel(os.path.join(local_path, i))
df = df.fillna("")
st.write(df)
csv = convert_df(df)
st.download_button(
label="Download data as csv",
data=csv,
file_name='large_df.csv',
mime='text/csv',
)