Upload 11 files

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/layout/layout1.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -10,12 +10,84 @@ from PIL import Image
 from rapid_layout import RapidLayout
 from rapid_orientation import RapidOrientation
 from rapid_table import RapidTable
+from streamlit_image_select import image_select
 
 orientation_engine = RapidOrientation()
 layout_engine = RapidLayout()
 table_engine = RapidTable()
 
 
+def orientation_part():
+    st.markdown('##### 文档图像方向分类')
+    img = image_select(
+        label="示例图（点击选择）",
+        images=[
+            'images/orientation/img_rot0_demo.jpg',
+            'images/orientation/rot_90.jpg',
+            'images/orientation/img_rot180_demo.jpg',
+            'images/orientation/rot_270.jpg'
+        ],
+        captions=['0', '90', '180', '270'],
+        key='cls'
+    )
+    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
+                                       key='orientation',
+                                       label_visibility='collapsed')
+
+    col1, col2 = st.columns([5, 5])
+    img_empty = col1.empty()
+
+    image, orientation_res, elapse = get_cls_res(img)
+    if img_file_buffer:
+        image, orientation_res, elapse = get_cls_res(img_file_buffer)
+
+    img_empty.image(image, use_column_width=True)
+    col2.markdown(f'- 方向分类结果：{orientation_res}° \n - 耗费时间：{elapse:.4f}s')
+
+
+def get_cls_res(img_file_buffer):
+    image = Image.open(img_file_buffer)
+    img = np.array(image)
+    img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    orientation_res, elapse = orientation_engine(img_array)
+    return image, orientation_res, elapse
+
+
+def layout_part():
+    st.markdown('##### 文档图像版面分析')
+    img = image_select(
+        label="示例图（点击选择）",
+        images=[
+            'images/layout/layout.png',
+            'images/layout/layout1.png',
+            'images/layout/layout2.png',
+        ],
+        key='layout_select'
+    )
+    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
+                                       key='layout',
+                                       label_visibility='collapsed')
+
+    layout_col1, layout_col2 = st.columns([5, 5])
+    img_empty = layout_col1.empty()
+
+    image, drawed_img = get_layout_res(img)
+    if img_file_buffer:
+        image, drawed_img = get_layout_res(img_file_buffer)
+
+    img_empty.image(image, use_column_width=True)
+    layout_col2.image(drawed_img, use_column_width=True)
+
+
+def get_layout_res(img_file_buffer):
+    image = Image.open(img_file_buffer)
+    img = np.array(image)
+    img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    layout_res, _ = layout_engine(img_array)
+    drawed_img = vis_layout(img, layout_res)
+    return image, drawed_img
+
+
 def vis_layout(img: np.ndarray, layout_res: list) -> None:
     tmp_img = copy.deepcopy(img)
     for v in layout_res:
@@ -25,12 +97,47 @@ def vis_layout(img: np.ndarray, layout_res: list) -> None:
         start_point = (bbox[0], bbox[1])
         end_point = (bbox[2], bbox[3])
 
-        cv2.rectangle(tmp_img, start_point, end_point, (0, 255, 0), 2)
+        cv2.rectangle(tmp_img, start_point, end_point, (0, 0, 0), 2)
         cv2.putText(tmp_img, label, start_point,
-                    cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
+                    cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 2)
     return tmp_img
 
 
+def table_part():
+    st.markdown('##### 表格还原')
+    img = image_select(
+        label="示例图（点击选择）",
+        images=[
+            'images/table/table.jpg',
+            'images/table/table1.png',
+            'images/table/table2.png',
+        ],
+        key='table_select'
+    )
+    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
+                                       key='table',
+                                       label_visibility='collapsed')
+
+    table_col1, table_col2 = st.columns([5, 5])
+    img_empty = table_col1.empty()
+
+    image, table_html_str = get_table_res(img)
+    if img_file_buffer:
+        image, table_html_str = get_table_res(img_file_buffer)
+
+    img_empty.image(image, use_column_width=True)
+    table_col2.markdown(table_html_str, unsafe_allow_html=True)
+
+
+def get_table_res(img_file_buffer):
+    image = Image.open(img_file_buffer)
+    img = np.array(image)
+    img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    table_html_str, _ = table_engine(img_array)
+    table_html_str = vis_table(table_html_str)
+    return image, table_html_str
+
+
 def vis_table(table_res) -> str:
     style_res = '''<style>td {border-left: 1px solid;border-bottom:1px solid;}
                    table, th {border-top:1px solid;font-size: 10px;
@@ -55,55 +162,6 @@ if __name__ == '__main__':
 
     img_suffix = ["png", "jpg", "jpeg"]
 
-    st.markdown('##### [文档图像方向分类](https://github.com/RapidAI/RapidStructure/blob/main/docs/README_Orientation.md)')
-
-    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
-                                       key='orientation',
-                                       label_visibility='collapsed')
-    col1, col2 = st.columns([5, 5])
-
-    img_empty = col1.empty()
-    if img_file_buffer:
-        image = Image.open(img_file_buffer)
-        img = np.array(image)
-        img_empty.image(image, use_column_width=True)
-
-        img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        orientation_res, elapse = orientation_engine(img)
-
-        col2.markdown(f'- 方向分类结果：{orientation_res}° \n - 耗费时间：{elapse:.4f}s')
-
-    st.markdown('##### [文档图像版面分析](https://github.com/RapidAI/RapidStructure/blob/main/docs/README_Layout.md)')
-    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
-                                       key='layout',
-                                       label_visibility='collapsed')
-    layout_col1, layout_col2 = st.columns([5, 5])
-
-    img_empty = layout_col1.empty()
-    if img_file_buffer:
-        image = Image.open(img_file_buffer)
-        img = np.array(image)
-        img_empty.image(image, use_column_width=True)
-
-        img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        layout_res, _ = layout_engine(img)
-        drawed_img = vis_layout(img, layout_res)
-        layout_col2.image(drawed_img, use_column_width=True)
-
-    st.markdown('##### [表格还原](https://github.com/RapidAI/RapidStructure/blob/main/docs/README_Table.md)')
-    img_file_buffer = st.file_uploader("Upload an image", type=img_suffix,
-                                       key='table',
-                                       label_visibility='collapsed')
-    table_col1, table_col2 = st.columns([5, 5])
-
-    img_empty = table_col1.empty()
-    if img_file_buffer:
-        image = Image.open(img_file_buffer)
-        img = np.array(image)
-        img_empty.image(image, use_column_width=True)
-
-        img_array = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        table_html_str, _ = table_engine(img)
-
-        table_html_str = vis_table(table_html_str)
-        table_col2.markdown(table_html_str, unsafe_allow_html=True)
+    orientation_part()
+    layout_part()
+    table_part()