from model import OCRModel, DET_ARCHS, RECO_ARCHS
from ui import OCRUI
from doctr.io import DocumentFile
from doctr.utils.visualization import visualize_page
import matplotlib.pyplot as plt
import time
import streamlit as st
import aiohttp
import asyncio
def main():
    ui = OCRUI()
    model = OCRModel()
    
    uploaded_file, params, current_page = ui.setup_sidebar(DET_ARCHS, RECO_ARCHS)
    
    if st.sidebar.button("分析步骤演示"):
        if uploaded_file is None:
            st.sidebar.error("请上传一个文档")
            return
            
        if current_page is None:
            st.sidebar.error("页面加载失败")
            return

        try:
            with st.spinner('正在处理中...请稍等'):
                st.sidebar.info("开始处理...")
                # Process page
                start_time = time.time()
                st.sidebar.info("加载模型...")
                model.load_model(**params)
                st.sidebar.info("分析页面...")
                seg_map, out = model.process_page(current_page)
                
                # Display results
                st.sidebar.info("显示结果...")
                fig, ax = plt.subplots()
                ax.imshow(seg_map)
                ax.axis("off")
                ui.cols[1].pyplot(fig)
                plt.close(fig)
                fig = visualize_page(out.pages[0].export(), out.pages[0].page, 
                                   interactive=False, add_labels=True)
                ui.cols[2].pyplot(fig)
                plt.close(fig)
                 # Reconstructed page
                 # Get reconstructed page
                reconstructed = model.get_reconstructed_page(out, current_page)
                ui.cols[3].image(reconstructed, clamp=True)
                plt.close(fig)
                
                # Display processing time and results
                total_time = time.time() - start_time
                st.success(f"处理完成! 用时: {total_time:.2f}秒")
                # Extract and display text results
                text_results = []
                for block in out.pages[0].export()["blocks"]:
                    for line in block["lines"]:
                        for word in line["words"]:
                            text_results.append(word["value"])
                
                st.write("识别结果:", " ".join(text_results))
                st.json({"total_time": total_time})
            
        except Exception as e:
            st.error(f"处理失败: {str(e)}")
            st.sidebar.error("发生错误，请查看详细信息")

if __name__ == "__main__":
    main()