Update app.py

app.py

@@ -5,6 +5,7 @@ from transformers import TrOCRProcessor, VisionEncoderDecoderModel, DonutProcess
 import torch
 import re
 import pytesseract
+from io import BytesIO
 
 
 def predict_arabic(img, model_name="UBC-NLP/Qalam"):
@@ -77,6 +78,8 @@ st.set_page_config(
 
 # Upload an image and set some options for demo purposes
 st.header("Qalam: A Multilingual OCR System")
+st.sidebar.header("Configuration and Image Upload")
+st.sidebar.subheader("Adjust Image Enhancement Options")
 img_file = st.sidebar.file_uploader(label='Upload a file', type=['png', 'jpg'])
 realtime_update = st.sidebar.checkbox(label="Update in Real Time", value=True)
 # box_color = st.sidebar.color_picker(label="Box Color", value='#0000FF')
@@ -86,6 +89,8 @@ aspect_dict = {
     "Free": None
 }
 aspect_ratio = aspect_dict[aspect_choice]
+st.sidebar.subheader("Select OCR Language and Model")
+
 Lng = st.sidebar.selectbox(label="Language", options=[
     "English", "Arabic", "French", "Korean", "Chinese"])
 
@@ -97,7 +102,9 @@ Models = {
     "Chinese": "Donut"
 }
  
-st.sidebar.write("# Model: ", Models[Lng])
+st.sidebar.markdown(f"### Selected Model: {Models[Lng]}")
+
+
 
 if img_file:
     img = Image.open(img_file)
@@ -106,7 +113,7 @@ if img_file:
 
     col1, col2 = st.columns(2)
     with col1:
-        st.header("Select Input Image")
+        st.subheader("Input: Upload and Crop Your Image")
     # Get a cropped image from the frontend
         cropped_img = st_cropper(
             img, 
@@ -118,24 +125,42 @@ if img_file:
         
     with col2:
     # Manipulate cropped image at will
-        st.header("Output Image")
+        st.subheader("Output: Preview and Analyze")
         # _ = cropped_img.thumbnail((150, 150))
         st.image(cropped_img)
         button = st.button("Run OCR")
         if button:
             with st.spinner('Running OCR...'):
                 if Lng == "Arabic":
-                    st.write("# Arabic Text:")
-                    st.write(predict_arabic(cropped_img))
+                    ocr_text = predict_arabic(cropped_img)
+                    st.subheader(f"OCR Results for {Lng}")
+                    st.write(ocr_text)
+                    text_file = BytesIO(ocr_text.encode())
+                    st.download_button('Download Text', text_file, file_name='ocr_text.txt')
                 elif Lng == "English":
-                    st.write("# English Text:")
-                    st.write(predict_english(cropped_img))
+                    ocr_text = predict_english(cropped_img)
+                    st.subheader(f"OCR Results for {Lng}")
+                    st.write(ocr_text)
+                    text_file = BytesIO(ocr_text.encode())
+                    st.download_button('Download Text', text_file, file_name='ocr_text.txt')
                 elif Lng == "French":
-                    st.write("# French Text:")
-                    st.write(predict_tesseract(cropped_img))
+                    ocr_text = predict_tesseract(cropped_img)
+                    st.subheader(f"OCR Results for {Lng}")
+                    st.write(ocr_text)
+                    text_file = BytesIO(ocr_text.encode())
+                    st.download_button('Download Text', text_file, file_name='ocr_text.txt')
                 elif Lng == "Korean":
-                    st.write("# Korean Text:")
-                    st.write(predict_english(cropped_img))
+                    ocr_text = predict_english(cropped_img)
+                    st.subheader(f"OCR Results for {Lng}")
+                    st.write(ocr_text)
+                    text_file = BytesIO(ocr_text.encode())
+                    st.download_button('Download Text', text_file, file_name='ocr_text.txt')
                 elif Lng == "Chinese":
-                    st.write("# Chinese Text:")
-                    st.write(predict_english(cropped_img))
+                    ocr_text = predict_english(cropped_img)
+                    st.subheader(f"OCR Results for {Lng}")
+                    st.write(ocr_text)
+                    text_file = BytesIO(ocr_text.encode())
+                    st.download_button('Download Text', text_file, file_name='ocr_text.txt')
+
+
+