Update app.py

app.py

@@ -1,70 +1,74 @@
-import tempfile
-import os
-
-# مسیر امن برای Hugging Face
-runtime_dir = os.path.join(tempfile.gettempdir(), ".streamlit")
-os.environ["STREAMLIT_RUNTIME_DIR"] = runtime_dir
-os.makedirs(runtime_dir, exist_ok=True)
+# app.py — TRUST OCR DEMO (Streamlit) for surya-ocr==0.4.14
 
-import argparse
+import os
 import io
+import tempfile
 from typing import List
 
+import numpy as np
+import cv2
+from PIL import Image
 import pypdfium2
+import pytesseract
 import streamlit as st
-# ✅ تغییر مهم
-from surya.ocr import batch_text_detection, run_ocr
+
+# ===== Safe runtime dir for Streamlit/HF cache (esp. in containers) =====
+runtime_dir = os.path.join(tempfile.gettempdir(), ".streamlit")
+os.environ["STREAMLIT_RUNTIME_DIR"] = runtime_dir
+os.makedirs(runtime_dir, exist_ok=True)
+
+# ===== Surya imports (v0.4.x) =====
+from surya.detection import batch_text_detection
 from surya.layout import batch_layout_detection
-from surya.model.detection.segformer import load_model, load_processor
+
+# Detection model loaders: prefer segformer; fallback to model (older path)
+try:
+    from surya.model.detection.segformer import load_model as load_det_model, load_processor as load_det_processor
+except ImportError:
+    from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
+
 from surya.model.recognition.model import load_model as load_rec_model
 from surya.model.recognition.processor import load_processor as load_rec_processor
-from surya.model.ordering.processor import load_processor as load_order_processor
+
 from surya.model.ordering.model import load_model as load_order_model
+from surya.model.ordering.processor import load_processor as load_order_processor
 from surya.ordering import batch_ordering
+
+from surya.ocr import run_ocr
 from surya.postprocessing.heatmap import draw_polys_on_image
 from surya.postprocessing.text import draw_text_on_image
-from PIL import Image
 from surya.languages import CODE_TO_LANGUAGE
 from surya.input.langs import replace_lang_with_code
 from surya.schema import OCRResult, TextDetectionResult, LayoutResult, OrderResult
-import pytesseract
-import cv2
-import numpy as np
 
-# -------------------
-# Args
-# -------------------
-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)
 
-try:
-    args = parser.parse_args()
-except SystemExit as e:
-    print(f"Error parsing arguments: {e}")
-    os._exit(e.code)
-
-# -------------------
-# Helper Functions
-# -------------------
-def remove_border(image_path, output_path):
+# ===================== Helper Functions =====================
+
+def remove_border(image_path: str, output_path: str) -> np.ndarray:
+    """Remove outer border & deskew (perspective) if a rectangular contour is found."""
     image = cv2.imread(image_path)
+    if image is None:
+        raise ValueError(f"Cannot read image: {image_path}")
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
     _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
     contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    max_contour = max(contours, key=cv2.contourArea)
+    if not contours:
+        cv2.imwrite(output_path, image)
+        return image
 
+    max_contour = max(contours, key=cv2.contourArea)
     epsilon = 0.02 * cv2.arcLength(max_contour, True)
     approx = cv2.approxPolyDP(max_contour, epsilon, True)
 
     if len(approx) == 4:
-        pts = approx.reshape(4, 2)
+        pts = approx.reshape(4, 2).astype("float32")
         rect = np.zeros((4, 2), dtype="float32")
         s = pts.sum(axis=1)
-        rect[0] = pts[np.argmin(s)]
-        rect[2] = pts[np.argmax(s)]
+        rect[0] = pts[np.argmin(s)]   # tl
+        rect[2] = pts[np.argmax(s)]   # br
         diff = np.diff(pts, axis=1)
-        rect[1] = pts[np.argmin(diff)]
-        rect[3] = pts[np.argmax(diff)]
+        rect[1] = pts[np.argmin(diff)]  # tr
+        rect[3] = pts[np.argmax(diff)]  # bl
         (tl, tr, br, bl) = rect
         widthA = np.linalg.norm(br - bl)
         widthB = np.linalg.norm(tr - tl)
@@ -73,8 +77,8 @@ def remove_border(image_path, output_path):
         heightB = np.linalg.norm(tl - bl)
         maxHeight = max(int(heightA), int(heightB))
         dst = np.array([[0, 0], [maxWidth - 1, 0],
-                       [maxWidth - 1, maxHeight - 1],
-                       [0, maxHeight - 1]], dtype="float32")
+                        [maxWidth - 1, maxHeight - 1],
+                        [0, maxHeight - 1]], dtype="float32")
         M = cv2.getPerspectiveTransform(rect, dst)
         cropped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
         cv2.imwrite(output_path, cropped)
@@ -84,82 +88,71 @@ def remove_border(image_path, output_path):
         return image
 
 
-def text_detection(img):
-    pred = batch_text_detection([img], det_model, det_processor)[0]
-    polygons = [p.polygon for p in pred.bboxes]
-    det_img = draw_polys_on_image(polygons, img.copy())
-    return det_img, pred
-
-
-def layout_detection(img):
-    _, det_pred = text_detection(img)
-    pred = batch_layout_detection([img], layout_model, layout_processor, [det_pred])[0]
-    polygons = [p.polygon for p in pred.bboxes]
-    labels = [p.label for p in pred.bboxes]
-    layout_img = draw_polys_on_image(polygons, img.copy(), labels=labels, label_font_size=40)
-    return layout_img, pred
-
-
-def order_detection(img):
-    _, layout_pred = layout_detection(img)
-    bboxes = [l.bbox for l in layout_pred.bboxes]
-    pred = batch_ordering([img], [bboxes], order_model, order_processor)[0]
-    polys = [l.polygon for l in pred.bboxes]
-    positions = [str(l.position) for l in pred.bboxes]
-    order_img = draw_polys_on_image(polys, img.copy(), labels=positions, label_font_size=40)
-    return order_img, pred
-
-
-def ocr(img, langs: List[str]):
-    replace_lang_with_code(langs)
-    img_pred = run_ocr([img], [langs], det_model, det_processor, rec_model, rec_processor)[0]
-    bboxes = [l.bbox for l in img_pred.text_lines]
-    text = [l.text for l in img_pred.text_lines]
-    rec_img = draw_text_on_image(bboxes, text, img.size, langs, has_math="_math" in langs)
-    return rec_img, img_pred
-
-
-def open_pdf(pdf_file):
+def open_pdf(pdf_file) -> pypdfium2.PdfDocument:
     stream = io.BytesIO(pdf_file.getvalue())
     return pypdfium2.PdfDocument(stream)
 
 
-@st.cache_data()
-def get_page_image(pdf_file, page_num, dpi=96):
+@st.cache_data(show_spinner=False)
+def get_page_image(pdf_file, page_num: int, dpi: int = 96) -> Image.Image:
     doc = open_pdf(pdf_file)
     renderer = doc.render(pypdfium2.PdfBitmap.to_pil, page_indices=[page_num - 1], scale=dpi / 72)
     png = list(renderer)[0]
     return png.convert("RGB")
 
 
-@st.cache_data()
-def page_count(pdf_file):
+@st.cache_data(show_spinner=False)
+def page_count(pdf_file) -> int:
     doc = open_pdf(pdf_file)
     return len(doc)
 
-# -------------------
-# Streamlit UI Config
-# -------------------
-st.set_page_config(layout="wide")
+
+# ===================== Streamlit UI =====================
+
+st.set_page_config(page_title="TRUST OCR DEMO", layout="wide")
+st.markdown("# TRUST OCR DEMO")
+
+# Sidebar controls
+in_file = st.sidebar.file_uploader("فایل PDF یا عکس :", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"])
+languages = st.sidebar.multiselect(
+    "زبان‌ها (Languages)",
+    sorted(list(CODE_TO_LANGUAGE.values())),
+    default=["Persian"],
+    max_selections=4
+)
+auto_rotate = st.sidebar.toggle("چرخش خودکار (Tesseract OSD)", value=True)
+auto_border = st.sidebar.toggle("حذف قاب/کادر تصویر ورودی", value=True)
+
+text_det_btn = st.sidebar.button("تشخیص متن (Detection)")
+layout_det_btn = st.sidebar.button("آنالیز صفحه (Layout)")
+order_det_btn = st.sidebar.button("ترتیب خوانش (Reading Order)")
+text_rec_btn = st.sidebar.button("تبدیل به متن (Recognition)")
+
+if in_file is None:
+    st.info("یک فایل PDF/عکس از سایدبار انتخاب کنید. | Please upload a file to begin.")
+    st.stop()
+
+filetype = in_file.type
+
+# Two-column layout (left: outputs / right: input image)
 col2, col1 = st.columns([.5, .5])
 
-# -------------------
-# Load Models
-# -------------------
-@st.cache_resource()
+# ===================== Load Models (cached) =====================
+
+@st.cache_resource(show_spinner=True)
 def load_det_cached():
-    return load_model(checkpoint="vikp/surya_det2"), load_processor(checkpoint="vikp/surya_det2")
+    return load_det_model(checkpoint="vikp/surya_det2"), load_det_processor(checkpoint="vikp/surya_det2")
 
-@st.cache_resource()
+@st.cache_resource(show_spinner=True)
 def load_rec_cached():
     return load_rec_model(checkpoint="MohammadReza-Halakoo/TrustOCR"), \
            load_rec_processor(checkpoint="MohammadReza-Halakoo/TrustOCR")
 
-@st.cache_resource()
+@st.cache_resource(show_spinner=True)
 def load_layout_cached():
-    return load_model(checkpoint="vikp/surya_layout2"), load_processor(checkpoint="vikp/surya_layout2")
+    return load_det_model(checkpoint="vikp/surya_layout2"), load_det_processor(checkpoint="vikp/surya_layout2")
 
-@st.cache_resource()
+@st.cache_resource(show_spinner=True)
 def load_order_cached():
     return load_order_model(checkpoint="vikp/surya_order"), load_order_processor(checkpoint="vikp/surya_order")
 
@@ -169,20 +162,74 @@ rec_model, rec_processor = load_rec_cached()
 layout_model, layout_processor = load_layout_cached()
 order_model, order_processor = load_order_cached()
 
-# -------------------
-# UI
-# -------------------
-st.markdown("# TRUST OCR DEMO")
 
-in_file = st.sidebar.file_uploader("فایل PDF یا عکس :", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"])
-languages = st.sidebar.multiselect("زبان‌ها", sorted(list(CODE_TO_LANGUAGE.values())), default=["Persian"], max_selections=4)
+# ===================== High-level Ops =====================
 
-if in_file is None:
-    st.stop()
+def _apply_auto_rotate(pil_img: Image.Image) -> Image.Image:
+    """Auto-rotate using Tesseract OSD if enabled."""
+    if not auto_rotate:
+        return pil_img
+    try:
+        osd = pytesseract.image_to_osd(pil_img, output_type=pytesseract.Output.DICT)
+        angle = int(osd.get("rotate", 0))  # 0/90/180/270
+        if angle and angle % 360 != 0:
+            # Tesseract returns counter-clockwise; PIL rotates counter-clockwise with positive values
+            return pil_img.rotate(-angle, expand=True)
+        return pil_img
+    except Exception as e:
+        st.warning(f"OSD rotation failed, continuing without rotation. Error: {e}")
+        return pil_img
+
+
+def text_detection(pil_img: Image.Image):
+    """Text block detection via Surya detection pipeline."""
+    pred: TextDetectionResult = batch_text_detection([pil_img], det_model, det_processor)[0]
+    polygons = [p.polygon for p in pred.bboxes]
+    det_img = draw_polys_on_image(polygons, pil_img.copy())
+    return det_img, pred
+
+
+def layout_detection(pil_img: Image.Image):
+    """Page layout analysis (requires detection result)."""
+    _, det_pred = text_detection(pil_img)
+    pred: LayoutResult = batch_layout_detection([pil_img], layout_model, layout_processor, [det_pred])[0]
+    polygons = [p.polygon for p in pred.bboxes]
+    labels = [p.label for p in pred.bboxes]
+    layout_img = draw_polys_on_image(polygons, pil_img.copy(), labels=labels, label_font_size=40)
+    return layout_img, pred
+
+
+def order_detection(pil_img: Image.Image):
+    """Reading order estimation (requires layout result)."""
+    _, layout_pred = layout_detection(pil_img)
+    bboxes = [l.bbox for l in layout_pred.bboxes]
+    pred: OrderResult = batch_ordering([pil_img], [bboxes], order_model, order_processor)[0]
+    polys = [l.polygon for l in pred.bboxes]
+    positions = [str(l.position) for l in pred.bboxes]
+    order_img = draw_polys_on_image(polys, pil_img.copy(), labels=positions, label_font_size=40)
+    return order_img, pred
+
+
+def ocr_page(pil_img: Image.Image, langs: List[str]):
+    """Full-page OCR using Surya run_ocr."""
+    # User selects languages by names; convert to codes
+    replace_lang_with_code(langs)  # in-place
+    img_pred: OCRResult = run_ocr([pil_img], [langs], det_model, det_processor, rec_model, rec_processor)[0]
+    bboxes = [l.bbox for l in img_pred.text_lines]
+    text = [l.text for l in img_pred.text_lines]
+    rec_img = draw_text_on_image(bboxes, text, pil_img.size, langs, has_math="_math" in langs)
+    return rec_img, img_pred
+
+
+# ===================== Input Handling =====================
 
-filetype = in_file.type
 if "pdf" in filetype:
-    page_number = st.sidebar.number_input(f"صفحه:", min_value=1, value=1, max_value=page_count(in_file))
+    try:
+        pg_cnt = page_count(in_file)
+    except Exception as e:
+        st.error(f"خواندن PDF ناموفق بود: {e}")
+        st.stop()
+    page_number = st.sidebar.number_input("صفحه:", min_value=1, value=1, max_value=pg_cnt)
     pil_image = get_page_image(in_file, page_number)
 else:
     bytes_data = in_file.getvalue()
@@ -191,40 +238,55 @@ else:
     file_path = os.path.join(temp_dir, in_file.name)
     with open(file_path, "wb") as f:
         f.write(bytes_data)
-    out_file = file_path.split('.')[0] + "-1.JPG"
-    remove_border(file_path, out_file)
-    pil_image = Image.open(out_file).convert("RGB")
-
-text_det = st.sidebar.button("تشخیص متن")
-text_rec = st.sidebar.button("تبدیل به متن")
-layout_det = st.sidebar.button("آنالیز صفحه")
-order_det = st.sidebar.button("ترتیب خوانش")
-
-if text_det:
-    osd = pytesseract.image_to_osd(pil_image, output_type='dict')
-    im_fixed = pil_image.copy().rotate(osd['orientation'])
-    det_img, pred = text_detection(im_fixed)
-    with col1:
-        st.image(det_img, caption="تشخیص متن", use_column_width=True)
-
-if layout_det:
-    layout_img, pred = layout_detection(pil_image)
-    with col1:
-        st.image(layout_img, caption="آنالیز صفحه", use_column_width=True)
-
-if text_rec:
-    rec_img, pred = ocr(pil_image, languages)
-    with col1:
-        text_tab, json_tab = st.tabs(["متن صفحه", "JSON"])
-        with text_tab:
-            st.text("\n".join([p.text for p in pred.text_lines]))
-        with json_tab:
-            st.json(pred.model_dump(), expanded=True)
-
-if order_det:
-    order_img, pred = order_detection(pil_image)
-    with col1:
-        st.image(order_img, caption="ترتیب خوانش", use_column_width=True)
+    out_file = os.path.splitext(file_path)[0] + "-1.JPG"
+    try:
+        if auto_border:
+            _ = remove_border(file_path, out_file)
+            pil_image = Image.open(out_file).convert("RGB")
+        else:
+            pil_image = Image.open(file_path).convert("RGB")
+    except Exception as e:
+        st.warning(f"حذف قاب/بازخوانی تصویر با خطا مواجه شد؛ تصویر اصلی استفاده می‌شود. Error: {e}")
+        pil_image = Image.open(file_path).convert("RGB")
+
+# Auto-rotate if enabled
+pil_image = _apply_auto_rotate(pil_image)
+
+# ===================== Buttons Logic =====================
+
+with col1:
+    if text_det_btn:
+        try:
+            det_img, det_pred = text_detection(pil_image)
+            st.image(det_img, caption="تشخیص متن (Detection)", use_column_width=True)
+        except Exception as e:
+            st.error(f"خطا در تشخیص متن: {e}")
+
+    if layout_det_btn:
+        try:
+            layout_img, layout_pred = layout_detection(pil_image)
+            st.image(layout_img, caption="آنالیز صفحه (Layout)", use_column_width=True)
+        except Exception as e:
+            st.error(f"خطا در آنالیز صفحه: {e}")
+
+    if order_det_btn:
+        try:
+            order_img, order_pred = order_detection(pil_image)
+            st.image(order_img, caption="ترتیب خوانش (Reading Order)", use_column_width=True)
+        except Exception as e:
+            st.error(f"خطا در ترتیب خوانش: {e}")
+
+    if text_rec_btn:
+        try:
+            lang_names = list(languages) if languages else ["Persian"]
+            rec_img, ocr_pred = ocr_page(pil_image, lang_names)
+            text_tab, json_tab = st.tabs(["متن صفحه | Page Text", "JSON"])
+            with text_tab:
+                st.text("\n".join([p.text for p in ocr_pred.text_lines]))
+            with json_tab:
+                st.json(ocr_pred.model_dump(), expanded=False)
+        except Exception as e:
+            st.error(f"خطا در بازشناسی متن (Recognition): {e}")
 
 with col2:
-    st.image(pil_image, caption="تصویر ورودی", use_column_width=True)
+    st.image(pil_image, caption="تصویر ورودی | Input Preview", use_column_width=True)