UI 分离

app.py

@@ -1,209 +1,42 @@
-# Copyright (C) 2021-2024, Mindee.
-
-# This program is licensed under the Apache License 2.0.
-# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
-
-import cv2
-import matplotlib.pyplot as plt
-import numpy as np
-import streamlit as st
-import time
-from doctr.file_utils import is_tf_available
+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 time
+import streamlit as st
 
-def setup_device():
-    """Setup and return compute device configuration"""
-    selected_device = "cpu"  # Default to CPU
-    if torch.cuda.is_available():
-        device_options = ["cuda", "cpu"]
-        selected_device = st.sidebar.selectbox("计算设备", device_options)
-        forward_device = torch.device("cuda:0" if selected_device == "cuda" else "cpu")
-        
-        # Display GPU info if CUDA selected
-        st.sidebar.markdown(f"**当前设备**: {forward_device}")
-        if selected_device == "cuda":
-            st.sidebar.markdown(f"**GPU型号**: {torch.cuda.get_device_name(0)}")
-            st.sidebar.markdown(f"**可用显存**: {torch.cuda.get_device_properties(0).total_memory/1024/1024:.0f}MB")
-    else:
-        st.sidebar.write("当前仅支持CPU")
-        forward_device = torch.device("cpu")
-        st.sidebar.markdown(f"**当前设备**: {forward_device}")
+def main():
+    ui = OCRUI()
+    model = OCRModel()
+    
+    uploaded_file, params = ui.setup_sidebar(DET_ARCHS, RECO_ARCHS)
     
-    return forward_device, selected_device
-
-def format_time(seconds):
-    """Format seconds into human readable string"""
-    return f"{seconds:.2f}秒"
-
-if is_tf_available():
-    import tensorflow as tf
-    from backend.tensorflow import DET_ARCHS, RECO_ARCHS, forward_image, load_predictor
-
-    if any(tf.config.experimental.list_physical_devices("gpu")):
-        forward_device = tf.device("/gpu:0")
-    else:
-        forward_device = tf.device("/cpu:0")
-
-else:
-    import torch
-    from backend.pytorch import DET_ARCHS, RECO_ARCHS, forward_image, load_predictor
-
-    forward_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-def main(det_archs, reco_archs):
-    """Build a streamlit layout"""
-    # Wide mode
-    st.set_page_config(layout="wide")
-
-    # Designing the interface
-    st.title("美宜家文档文本识别DEMO")
-    # For newline
-    st.write("\n")
-    # Instructions
-    st.markdown("*提示：单击图像的右上角可以放大！*")
-    # Set the columns
-    cols = st.columns((1, 1, 1, 1))
-    cols[0].subheader("输入页面")
-    cols[1].subheader("分割热图")
-    cols[2].subheader("OCR 输出")
-    cols[3].subheader("页面重构")
-
-    # Sidebar
-    # File selection
-    st.sidebar.title("文档选择")
-    # Choose your own image
-    uploaded_file = st.sidebar.file_uploader("上传文件", type=["pdf", "png", "jpeg", "jpg"])
-    if uploaded_file is not None:
-        if uploaded_file.name.endswith(".pdf"):
-            doc = DocumentFile.from_pdf(uploaded_file.read())
-        else:
-            doc = DocumentFile.from_images(uploaded_file.read())
-        page_idx = st.sidebar.selectbox("页面选择", [idx + 1 for idx in range(len(doc))]) - 1
-        page = doc[page_idx]
-        cols[0].image(page)
-    # Hardware selection    
-    st.sidebar.title("硬件选择")
-    forward_device, selected_device = setup_device()
-    # Model selection
-    st.sidebar.title("模型选择")
-    st.sidebar.markdown("**后端**: " + ("TensorFlow" if is_tf_available() else "PyTorch"))
-    det_arch = st.sidebar.selectbox("文本检测模型", det_archs)
-    reco_arch = st.sidebar.selectbox("文本识别模型", reco_archs)
-
-    # For newline
-    st.sidebar.write("\n")
-    # Only straight pages or possible rotation
-    st.sidebar.title("参数")
-    assume_straight_pages = st.sidebar.checkbox("假设页面是直的", value=True)
-    # Disable page orientation detection
-    disable_page_orientation = st.sidebar.checkbox("禁用页面方向检测", value=False)
-    # Disable crop orientation detection
-    disable_crop_orientation = st.sidebar.checkbox("禁用裁剪方向检测", value=False)
-    # Straighten pages
-    straighten_pages = st.sidebar.checkbox("矫正页面", value=False)
-    # Export as straight boxes
-    export_straight_boxes = st.sidebar.checkbox("导出为直边框", value=False)
-    st.sidebar.write("\n")
-    # Binarization threshold
-    bin_thresh = st.sidebar.slider("二值化阈值", min_value=0.1, max_value=0.9, value=0.3, step=0.1)
-    st.sidebar.write("\n")
-    # Box threshold
-    box_thresh = st.sidebar.slider("边框阈值", min_value=0.1, max_value=0.9, value=0.1, step=0.1)
-    st.sidebar.write("\n")
-
     if st.sidebar.button("分析页面"):
         if uploaded_file is None:
             st.sidebar.write("请上传一个文档")
+            return
 
-        else:
-            start_model = time.time()
-            with st.spinner("加载模型..."):
-                predictor = load_predictor(
-                    det_arch=det_arch,
-                    reco_arch=reco_arch,
-                    assume_straight_pages=assume_straight_pages,
-                    straighten_pages=straighten_pages,
-                    export_as_straight_boxes=export_straight_boxes,
-                    disable_page_orientation=disable_page_orientation,
-                    disable_crop_orientation=disable_crop_orientation,
-                    bin_thresh=bin_thresh,
-                    box_thresh=box_thresh,
-                    device=forward_device,
-                )
-            model_time = time.time() - start_model
-            with st.spinner("分析中..."):
-                seg_time_start = time.time()
-                # Forward the image to the model
-                seg_map = forward_image(predictor, page, forward_device)
-                seg_map = np.squeeze(seg_map)
-                seg_map = cv2.resize(seg_map, (page.shape[1], page.shape[0]), interpolation=cv2.INTER_LINEAR)
-                seg_time = time.time() - seg_time_start
-                # Plot the raw heatmap
-                fig, ax = plt.subplots()
-                ax.imshow(seg_map)
-                ax.axis("off")
-                cols[1].pyplot(fig)
-
-                # Plot OCR output
-                ocr_time_start = time.time()
-                out = predictor([page])
-                fig = visualize_page(out.pages[0].export(), out.pages[0].page, interactive=False, add_labels=True)  # 改为True显示标签
-                cols[2].pyplot(fig)
-                ocr_time = time.time() - ocr_time_start
-
-               # Page reconsitution under input page
-                page_time_start = time.time()
-                page_export = out.pages[0].export()
-                if assume_straight_pages or (not assume_straight_pages and straighten_pages):
-                    # 获取合成图像
-                    img = out.pages[0].synthesize()
-                    
-                    # 计算所有文本框的边界
-                    x_min, y_min = float('inf'), float('inf')
-                    x_max, y_max = 0, 0
-                    
-                    for block in page_export["blocks"]:
-                        # 获取每个块的坐标
-                        coords = np.array(block["geometry"])
-                        x_min = min(x_min, coords[:, 0].min() * img.shape[1])
-                        y_min = min(y_min, coords[:, 1].min() * img.shape[0])
-                        x_max = max(x_max, coords[:, 0].max() * img.shape[1])
-                        y_max = max(y_max, coords[:, 1].max() * img.shape[0])
-    
-                    # 添加边距
-                    margin = 10
-                    x_min = max(0, x_min - margin)
-                    y_min = max(0, y_min - margin)
-                    x_max = min(img.shape[1], x_max + margin)
-                    y_max = min(img.shape[0], y_max + margin)
-                    
-                    # 裁剪图像
-                    cropped_img = img[int(y_min):int(y_max), int(x_min):int(x_max)]
-                    
-                    # 显示裁剪后的图像
-                    cols[3].image(cropped_img, clamp=True)
-
-                # 添加文本结果显示
-                page_time= time.time() - page_time_start
-                
-                total_time = time.time() - seg_time_start
-                
-                cols[0].subheader(f"输入页面 (总耗时: {format_time(total_time)})")
-                cols[1].subheader(f"分割热图 (耗时: {format_time(seg_time)})")
-                cols[2].subheader(f"OCR输出 (耗时: {format_time(ocr_time)})")
-                cols[3].subheader(f"页面重构 (模型加载: {format_time(page_time)})")
-                st.markdown("\n### OCR Text Results:")
-                for block in page_export["blocks"]:
-                    for line in block["lines"]:
-                        for word in line["words"]:
-                            st.write(f"Text: {word['value']}, Confidence: {word['confidence']:.2f}")
-
-                # Display JSON
-                st.markdown("\nHere are your analysis results in JSON format:")
-                #show total_time
-                st.json({"total_time": total_time}, expanded=True)
-                st.json(page_export, expanded=True)  # 改为True展开显示
+        doc = DocumentFile.from_pdf(uploaded_file.read()) if uploaded_file.name.endswith(".pdf") else DocumentFile.from_images(uploaded_file.read())
+        page_idx = st.sidebar.selectbox("页面选择", [idx + 1 for idx in range(len(doc))]) - 1
+        page = doc[page_idx]
+        
+        # Process page
+        start_time = time.time()
+        model.load_model(**params)
+        seg_map, out = model.process_page(page)
+        
+        # Display results
+        ui.cols[0].image(page)
+        fig, ax = plt.subplots()
+        ax.imshow(seg_map)
+        ax.axis("off")
+        ui.cols[1].pyplot(fig)
+        
+        fig = visualize_page(out.pages[0].export(), out.pages[0].page, interactive=False, add_labels=True)
+        ui.cols[2].pyplot(fig)
+        
+        # Display processing time and results
+        total_time = time.time() - start_time
+        st.json({"total_time": total_time, "results": out.pages[0].export()})
 
 if __name__ == "__main__":
-    main(DET_ARCHS, RECO_ARCHS)
+    main()

model.py

@@ -0,0 +1,39 @@
+import time
+from doctr.file_utils import is_tf_available
+import numpy as np
+import cv2
+
+if is_tf_available():
+    import tensorflow as tf
+    from backend.tensorflow import DET_ARCHS, RECO_ARCHS, forward_image, load_predictor
+else:
+    import torch
+    from backend.pytorch import DET_ARCHS, RECO_ARCHS, forward_image, load_predictor
+
+class OCRModel:
+    def __init__(self):
+        self.predictor = None
+        self.device = self._setup_device()
+
+    def _setup_device(self):
+        if is_tf_available():
+            if any(tf.config.experimental.list_physical_devices("gpu")):
+                return tf.device("/gpu:0")
+            return tf.device("/cpu:0")
+        else:
+            return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    def load_model(self, det_arch, reco_arch, **kwargs):
+        self.predictor = load_predictor(
+            det_arch=det_arch,
+            reco_arch=reco_arch,
+            device=self.device,
+            **kwargs
+        )
+
+    def process_page(self, page):
+        seg_map = forward_image(self.predictor, page, self.device)
+        seg_map = np.squeeze(seg_map)
+        seg_map = cv2.resize(seg_map, (page.shape[1], page.shape[0]), interpolation=cv2.INTER_LINEAR)
+        out = self.predictor([page])
+        return seg_map, out

ui.py

@@ -0,0 +1,40 @@
+import streamlit as st
+import matplotlib.pyplot as plt
+from doctr.utils.visualization import visualize_page
+
+class OCRUI:
+    def __init__(self):
+        self.setup_page_config()
+        self.cols = self.create_layout()
+
+    def setup_page_config(self):
+        st.set_page_config(layout="wide")
+        st.title("美宜家文档文本识别DEMO")
+        st.write("\n")
+        st.markdown("*提示：单击图像的右上角可以放大！*")
+
+    def create_layout(self):
+        cols = st.columns((1, 1, 1, 1))
+        cols[0].subheader("输入页面")
+        cols[1].subheader("分割热图")
+        cols[2].subheader("OCR 输出")
+        cols[3].subheader("页面重构")
+        return cols
+
+    def setup_sidebar(self, det_archs, reco_archs):
+        st.sidebar.title("文档选择")
+        uploaded_file = st.sidebar.file_uploader("上传文件", type=["pdf", "png", "jpeg", "jpg"])
+        
+        params = {
+            "assume_straight_pages": st.sidebar.checkbox("假设页面是直的", value=True),
+            "disable_page_orientation": st.sidebar.checkbox("禁用页面方向检测", value=False),
+            "disable_crop_orientation": st.sidebar.checkbox("禁用裁剪方向检测", value=False),
+            "straighten_pages": st.sidebar.checkbox("矫正页面", value=False),
+            "export_straight_boxes": st.sidebar.checkbox("导出为直边框", value=False),
+            "bin_thresh": st.sidebar.slider("二值化阈值", 0.1, 0.9, 0.3, 0.1),
+            "box_thresh": st.sidebar.slider("边框阈值", 0.1, 0.9, 0.1, 0.1),
+            "det_arch": st.sidebar.selectbox("文本检测模型", det_archs),
+            "reco_arch": st.sidebar.selectbox("文本识别模型", reco_archs)
+        }
+        
+        return uploaded_file, params