Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import streamlit as st
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
-import os
 import base64
 import uuid
 import time
@@ -12,10 +11,9 @@ from pathlib import Path
 os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 
 # Load tokenizer and model on CPU
-
 tokenizer = AutoTokenizer.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True)
 model = AutoModel.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-model = model.eval()
+model.eval()
 
 # Define folders for uploads and results
 UPLOAD_FOLDER = "./uploads"
@@ -80,37 +78,43 @@ st.set_page_config(page_title="GOT-OCR-2.0 Demo", layout="wide")
 
 uploaded_image = st.file_uploader("Upload your image", type=["png", "jpg", "jpeg"])
 
+# Create two columns for layout
+col1, col2 = st.columns(2)
+
 if uploaded_image:
     image = Image.open(uploaded_image)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
-
-    got_mode = st.selectbox("Choose one mode of GOT", [
-        "plain texts OCR",
-        "format texts OCR",
-        "plain multi-crop OCR",
-        "format multi-crop OCR",
-        "plain fine-grained OCR",
-        "format fine-grained OCR",
-    ])
-
-    fine_grained_mode = None
-    ocr_color = ""
-    ocr_box = ""
-
-    if "fine-grained" in got_mode:
-        fine_grained_mode = st.selectbox("Fine-grained type", ["box", "color"])
-        if fine_grained_mode == "box":
-            ocr_box = st.text_input("Input box: [x1,y1,x2,y2]", value="[0,0,100,100]")
-        elif fine_grained_mode == "color":
-            ocr_color = st.selectbox("Color list", ["red", "green", "blue"])
-
-    if st.button("Submit"):
-        with st.spinner("Processing..."):
-            result_text, html_result = run_GOT(image, got_mode, fine_grained_mode, ocr_color, ocr_box)
-            st.text_area("GOT Output", result_text, height=200)
-
-            if html_result:
-                st.markdown(html_result, unsafe_allow_html=True)
+    
+    with col1:
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+
+    with col2:
+        got_mode = st.selectbox("Choose one mode of GOT", [
+            "plain texts OCR",
+            "format texts OCR",
+            "plain multi-crop OCR",
+            "format multi-crop OCR",
+            "plain fine-grained OCR",
+            "format fine-grained OCR",
+        ])
+
+        fine_grained_mode = None
+        ocr_color = ""
+        ocr_box = ""
+
+        if "fine-grained" in got_mode:
+            fine_grained_mode = st.selectbox("Fine-grained type", ["box", "color"])
+            if fine_grained_mode == "box":
+                ocr_box = st.text_input("Input box: [x1,y1,x2,y2]", value="[0,0,100,100]")
+            elif fine_grained_mode == "color":
+                ocr_color = st.selectbox("Color list", ["red", "green", "blue"])
+
+        if st.button("Submit"):
+            with st.spinner("Processing..."):
+                result_text, html_result = run_GOT(image, got_mode, fine_grained_mode, ocr_color, ocr_box)
+                st.text_area("GOT Output", result_text, height=200)
+
+                if html_result:
+                    st.markdown(html_result, unsafe_allow_html=True)
 
 # Cleanup old files
 cleanup_old_files()