Update app.py

app.py

@@ -1,12 +1,14 @@
-
 import os
+os.environ["STREAMLIT_RUNTIME_DIR"] = "/app/.streamlit"
+os.makedirs("/app/.streamlit", exist_ok=True)
+
 import argparse
 import io
 from typing import List
 
 import pypdfium2
 import streamlit as st
-from surya.detection import batch_text_detection
+from surya.detection.infer import batch_text_detection   # اصلاح مسیر
 from surya.layout import batch_layout_detection
 from surya.model.detection.segformer import load_model, load_processor
 from surya.model.recognition.model import load_model as load_rec_model
@@ -25,53 +27,18 @@ import pytesseract
 import cv2
 import numpy as np
 
-from huggingface_hub import login
-
-os.environ["STREAMLIT_RUNTIME_DIR"] = "/app/.streamlit"
-os.makedirs("/app/.streamlit", exist_ok=True)
-
-
-
-
-# خواندن توکن از متغیر محیطی
-
-
-if "HUGGINGFACE_HUB_TOKEN" in os.environ:
-    login(token=os.environ["HUGGINGFACE_HUB_TOKEN"])
-
-
-# -------------------
-# Streamlit UI Config
-# -------------------
-st.set_page_config(layout="wide")
-col2, col1 = st.columns([.5, .5])
 
 # -------------------
-# Load Models
+# Args
 # -------------------
-@st.cache_resource()
-def load_det_cached():
-    return load_model(checkpoint="vikp/surya_det2"), load_processor(checkpoint="vikp/surya_det2")
+parser = argparse.ArgumentParser(description="Run OCR on an image or PDF.")
+parser.add_argument("--math", action="store_true", help="Use math model for detection", default=False)
 
-@st.cache_resource()
-def load_rec_cached():
-    # 🔥 اینجا مدل خودت که روی HF آپلود کردی
-    return load_rec_model(checkpoint="MohammadReza-Halakoo/TrustOCR"), \
-           load_rec_processor(checkpoint="MohammadReza-Halakoo/TrustOCR")
-
-@st.cache_resource()
-def load_layout_cached():
-    return load_model(checkpoint="vikp/surya_layout2"), load_processor(checkpoint="vikp/surya_layout2")
-
-@st.cache_resource()
-def load_order_cached():
-    return load_order_model(checkpoint="vikp/surya_order"), load_order_processor(checkpoint="vikp/surya_order")
-
-
-det_model, det_processor = load_det_cached()
-rec_model, rec_processor = load_rec_cached()
-layout_model, layout_processor = load_layout_cached()
-order_model, order_processor = load_order_cached()
+try:
+    args = parser.parse_args()
+except SystemExit as e:
+    print(f"Error parsing arguments: {e}")
+    os._exit(e.code)
 
 
 # -------------------
@@ -169,6 +136,39 @@ def page_count(pdf_file):
     return len(doc)
 
 
+# -------------------
+# Streamlit UI Config
+# -------------------
+st.set_page_config(layout="wide")
+col2, col1 = st.columns([.5, .5])
+
+# -------------------
+# Load Models
+# -------------------
+@st.cache_resource()
+def load_det_cached():
+    return load_model(checkpoint="vikp/surya_det2"), load_processor(checkpoint="vikp/surya_det2")
+
+@st.cache_resource()
+def load_rec_cached():
+    return load_rec_model(checkpoint="MohammadReza-Halakoo/TrustOCR"), \
+           load_rec_processor(checkpoint="MohammadReza-Halakoo/TrustOCR")
+
+@st.cache_resource()
+def load_layout_cached():
+    return load_model(checkpoint="vikp/surya_layout2"), load_processor(checkpoint="vikp/surya_layout2")
+
+@st.cache_resource()
+def load_order_cached():
+    return load_order_model(checkpoint="vikp/surya_order"), load_order_processor(checkpoint="vikp/surya_order")
+
+
+det_model, det_processor = load_det_cached()
+rec_model, rec_processor = load_rec_cached()
+layout_model, layout_processor = load_layout_cached()
+order_model, order_processor = load_order_cached()
+
+
 # -------------------
 # UI
 # -------------------