app.py

import os
import tempfile
import argparse
import io
from typing import List

import pypdfium2
import streamlit as st
from surya.ocr import run_ocr, 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
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.ordering import batch_ordering
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

# -------------------
# مسیر امن برای Streamlit در Hugging Face
# -------------------
runtime_dir = os.path.join(tempfile.gettempdir(), ".streamlit")
os.environ["STREAMLIT_RUNTIME_DIR"] = runtime_dir
os.makedirs(runtime_dir, exist_ok=True)

# -------------------
# 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):
    image = cv2.imread(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)

    epsilon = 0.02 * cv2.arcLength(max_contour, True)
    approx = cv2.approxPolyDP(max_contour, epsilon, True)

    if len(approx) == 4:
        pts = approx.reshape(4, 2)
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]
        rect[2] = pts[np.argmax(s)]
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]
        rect[3] = pts[np.argmax(diff)]
        (tl, tr, br, bl) = rect
        widthA = np.linalg.norm(br - bl)
        widthB = np.linalg.norm(tr - tl)
        maxWidth = max(int(widthA), int(widthB))
        heightA = np.linalg.norm(tr - br)
        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")
        M = cv2.getPerspectiveTransform(rect, dst)
        cropped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
        cv2.imwrite(output_path, cropped)
        return cropped
    else:
        cv2.imwrite(output_path, image)
        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):
    stream = io.BytesIO(pdf_file.getvalue())
    return pypdfium2.PdfDocument(stream)

@st.cache_data()
def get_page_image(pdf_file, page_num, dpi=96):
    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):
    doc = open_pdf(pdf_file)
    return len(doc)

# -------------------
# Streamlit UI
# -------------------
st.set_page_config(layout="wide")
col2, col1 = st.columns([.5, .5])

@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()

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)

if in_file is None:
    st.stop()

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))
    pil_image = get_page_image(in_file, page_number)
else:
    bytes_data = in_file.getvalue()
    temp_dir = "temp_files"
    os.makedirs(temp_dir, exist_ok=True)
    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)

with col2:
    st.image(pil_image, caption="تصویر ورودی", use_column_width=True)