Change layout

app.py

@@ -369,15 +369,13 @@ def main():
     st.title('Table Extraction Demo')
     st.write('\n')
 
-    cols = st.columns((1, 1))
-    cols[0].subheader('Input page')
-    cols[1].subheader('Table(s) detected')
+    tabs = st.tabs(
+        ['Table Detection', 'Table Structure Recognition']
+    )
 
-    st.sidebar.title('Image upload')
-    st.set_option('deprecation.showfileUploaderEncoding', False)
-    filename = st.sidebar.file_uploader('Upload files', type=['png', 'jpeg', 'jpg'])
+    filename = st.file_uploader('Upload image', type=['png', 'jpeg', 'jpg'])
 
-    if st.sidebar.button('Analyze image'):
+    if st.button('Analyze image'):
 
         if filename is None:
             st.sidebar.write('Please upload an image')
@@ -385,40 +383,48 @@ def main():
         else:
             print(filename)
             pil_img = PIL.Image.open(filename)
-            cols[0].image(pil_img)
 
             detection_result = table_detection(pil_img)
             crop_images, vis_det_img = crop_image(pil_img, detection_result)
-            cols[1].image(vis_det_img)
-
-            str_cols = st.columns((len(crop_images), ) * 5)
-            str_cols[0].subheader('Table image')
-            str_cols[1].subheader('OCR result')
-            str_cols[2].subheader('Structure result')
-            str_cols[3].subheader('Cells result')
-            str_cols[4].subheader('CSV result')
-
-            for i, img in enumerate(crop_images):
-                ocr_result = ocr(img)
-                structure_result = table_structure(img)
-                table_structures, cells, confidence_score = convert_stucture(ocr_result, img, structure_result)
-                cells = extract_text_from_cells(cells)
-                html_result = cells_to_html(cells)
-                df, csv_result = cells_to_csv(cells)
-                #print(df)
-
-                vis_ocr_img = visualize_ocr(img, ocr_result)
-                vis_str_img = visualize_structure(img, structure_result)
-                vis_cells_img = visualize_cells(img, cells)
-
-                str_cols[0].image(img)
-                str_cols[1].image(vis_ocr_img)
-                str_cols[2].image(vis_str_img)
-                str_cols[3].image(vis_cells_img)
-                #str_cols[4].dataframe(df)
-                str_cols[4].download_button('Download table', csv_result, f'table-{i}.csv', 'text/csv', key=f'download-csv-{i}')
-
-                st.markdown(html_result, unsafe_allow_html=True)
+
+            with tabs[0]:
+                st.image(vis_det_img)
+
+            with tabs[1]:
+                str_cols = st.columns((len(crop_images), ) * 5)
+                str_cols[0].subheader('Table image')
+                str_cols[1].subheader('OCR result')
+                str_cols[2].subheader('Structure result')
+                str_cols[3].subheader('Cells result')
+                str_cols[4].subheader('CSV result')
+
+                for i, img in enumerate(crop_images):
+                    ocr_result = ocr(img)
+                    structure_result = table_structure(img)
+                    table_structures, cells, confidence_score = convert_stucture(ocr_result, img, structure_result)
+                    cells = extract_text_from_cells(cells)
+                    html_result = cells_to_html(cells)
+                    df, csv_result = cells_to_csv(cells)
+                    #print(df)
+
+                    vis_ocr_img = visualize_ocr(img, ocr_result)
+                    vis_str_img = visualize_structure(img, structure_result)
+                    vis_cells_img = visualize_cells(img, cells)
+
+                    str_cols[0].image(img)
+                    str_cols[1].image(vis_ocr_img)
+                    str_cols[2].image(vis_str_img)
+                    str_cols[3].image(vis_cells_img)
+
+                    try:
+                        str_cols[4].dataframe(df)
+                    except:
+                        pass
+
+                    str_cols[4].download_button('Download table', csv_result, f'table-{i}.csv', 'text/csv', key=f'download-csv-{i}')
+
+                    st.write('\n')
+                    st.markdown(html_result, unsafe_allow_html=True)
 
 
 if __name__ == '__main__':