app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Sapphire OCR Engine V9 - Stable EasyOCR + Handwritten Arabic (TrOCR) + Tables
"""

import os
import gc
import time
import shutil
import logging
import threading
from PIL import Image

import cv2
import torch
import numpy as np
import gradio as gr
from huggingface_hub import HfApi, hf_hub_download
import fitz

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("SapphireOCR")

UPLOAD_FOLDER = "upload_tmp"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

REPO_ID = "Asem75/aiocr_asistant"
RAW_DIR = "raw_files"
PROCESS_DIR = "process_files"

# الإعدادات الافتراضية (خفيفة)
ENABLE_TABLE_DEFAULT = False
ENABLE_HEAVY_PREP_DEFAULT = False
DEFAULT_DPI = 200

ocr_stats = {
    "status": "🟢 المحرك جاهز (EasyOCR + يدوي عربي)",
    "processed_files_count": 0,
    "total_extracted_words": 0,
    "last_processed_file": "لا طلبات",
    "micro_logs": "🚀 الرادار يراقب الملفات..."
}

# ------------------------------
# EasyOCR (تحميل بطيء)
# ------------------------------
_easyocr_reader = None
def get_easyocr_reader():
    global _easyocr_reader
    if _easyocr_reader is None:
        import easyocr
        logger.info("📥 تحميل EasyOCR (عربي + إنجليزي)...")
        _easyocr_reader = easyocr.Reader(['ar', 'en'], gpu=torch.cuda.is_available())
    return _easyocr_reader

# ------------------------------
# نموذج الكتابة اليدوية العربية (TrOCR)
# ------------------------------
_handwritten_processor = None
_handwritten_model = None

def get_arabic_handwritten_model():
    global _handwritten_processor, _handwritten_model
    if _handwritten_model is None:
        from transformers import TrOCRProcessor, VisionEncoderDecoderModel
        model_name = "goforit18/trocr-base-arabic-handwritten"
        logger.info("📥 تحميل نموذج الكتابة اليدوية العربية...")
        _handwritten_processor = TrOCRProcessor.from_pretrained(model_name)
        _handwritten_model = VisionEncoderDecoderModel.from_pretrained(model_name)
        if torch.cuda.is_available():
            _handwritten_model.to("cuda")
    return _handwritten_processor, _handwritten_model

def extract_arabic_handwriting(image_pil: Image.Image) -> str:
    """
    استخراج النص العربي اليدوي من صورة PIL (RGB) باستخدام TrOCR.
    """
    processor, model = get_arabic_handwritten_model()
    if image_pil.mode != "RGB":
        image_pil = image_pil.convert("RGB")
    pixel_values = processor(images=image_pil, return_tensors="pt").pixel_values
    if torch.cuda.is_available():
        pixel_values = pixel_values.cuda()
    generated_ids = model.generate(pixel_values)
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return text

# ------------------------------
# Table Transformer (تحميل بطيء)
# ------------------------------
_table_processor = None
_table_detector = None
_table_structure_detector = None
def get_table_models():
    global _table_processor, _table_detector, _table_structure_detector
    if _table_detector is None:
        from transformers import AutoImageProcessor, TableTransformerForObjectDetection
        logger.info("📥 تحميل Table Transformer...")
        _table_processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-detection")
        _table_detector = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection")
        _table_structure_detector = TableTransformerForObjectDetection.from_pretrained(
            "microsoft/table-transformer-structure-recognition"
        )
        if torch.cuda.is_available():
            _table_detector.to("cuda")
            _table_structure_detector.to("cuda")
    return _table_processor, _table_detector, _table_structure_detector

# ------------------------------
# دوال المعالجة المسبقة (بدون تغيير)
# ------------------------------
def remove_shadows(img_gray, kernel_size=151):
    background = cv2.medianBlur(img_gray, kernel_size)
    diff = cv2.absdiff(img_gray, background)
    normalized = cv2.normalize(diff, None, 0, 255, cv2.NORM_MINMAX)
    return 255 - normalized

def correct_perspective(img_gray):
    _, thresh = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours: return img_gray
    largest = max(contours, key=cv2.contourArea)
    peri = cv2.arcLength(largest, True)
    approx = cv2.approxPolyDP(largest, 0.02 * peri, True)
    if len(approx) != 4: return img_gray
    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)]
    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.sqrt(((br[0]-bl[0])**2)+((br[1]-bl[1])**2))
    widthB = np.sqrt(((tr[0]-tl[0])**2)+((tr[1]-tl[1])**2))
    maxWidth = max(int(widthA), int(widthB))
    heightA = np.sqrt(((tr[0]-br[0])**2)+((tr[1]-br[1])**2))
    heightB = np.sqrt(((tl[0]-bl[0])**2)+((tl[1]-bl[1])**2))
    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)
    return cv2.warpPerspective(img_gray, M, (maxWidth, maxHeight))

def deskew_image(img_gray):
    _, binary = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    coords = np.column_stack(np.where(binary > 0))
    if len(coords) < 100: return img_gray
    rect = cv2.minAreaRect(coords)
    angle = rect[-1]
    if angle < -45: angle = -(90 + angle)
    else: angle = -angle
    if abs(angle) < 0.3: return img_gray
    (h, w) = img_gray.shape[:2]
    M = cv2.getRotationMatrix2D((w//2, h//2), angle, 1.0)
    return cv2.warpAffine(img_gray, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)

def smart_crop(img_gray, margin=20):
    _, binary = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    points = cv2.findNonZero(binary)
    if points is None: return img_gray
    x,y,w,h = cv2.boundingRect(points)
    x = max(x-margin, 0)
    y = max(y-margin, 0)
    w = min(w+2*margin, img_gray.shape[1]-x)
    h = min(h+2*margin, img_gray.shape[0]-y)
    return img_gray[y:y+h, x:x+w]

def illumination_correct(img_gray):
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    return clahe.apply(img_gray)

def enhance_arabic_edges(img_gray):
    blurred = cv2.GaussianBlur(img_gray, (0,0), 3)
    sharpened = cv2.addWeighted(img_gray, 2.0, blurred, -1.0, 0)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
    return cv2.morphologyEx(sharpened, cv2.MORPH_CLOSE, kernel)

def preprocess_image(img_np, heavy=False):
    if len(img_np.shape) == 3:
        gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
    else:
        gray = img_np.copy()
    if heavy:
        gray = remove_shadows(gray)
        gray = correct_perspective(gray)
    gray = smart_crop(gray)
    gray = deskew_image(gray)
    gray = smart_crop(gray)
    gray = illumination_correct(gray)
    gray = enhance_arabic_edges(gray)
    return gray   # صورة رمادية محسّنة

# ------------------------------
# دوال OCR المركزية (تختار بين EasyOCR و اليدوي)
# ------------------------------
def ocr_with_easyocr(image_np):
    reader = get_easyocr_reader()
    return " ".join(reader.readtext(image_np, detail=0))

def ocr_image(image_np_or_pil, handwritten=False):
    """
    إذا كان `handwritten` == True، نستخدم النموذج اليدوي (صورة PIL RGB).
    وإلا نستخدم EasyOCR (مصفوفة رمادية).
    """
    if handwritten:
        try:
            if isinstance(image_np_or_pil, np.ndarray):
                # تحويل المصفوفة الرمادية إلى PIL RGB
                pil_img = Image.fromarray(image_np_or_pil).convert("RGB")
            else:
                pil_img = image_np_or_pil.convert("RGB")
            return extract_arabic_handwriting(pil_img)
        except Exception as e:
            logger.warning(f"فشل استخراج الكتابة اليدوية، الرجوع إلى EasyOCR: {e}")
            # الرجوع إلى EasyOCR مع الصورة الرمادية
            if isinstance(image_np_or_pil, np.ndarray):
                gray = image_np_or_pil
            else:
                gray = np.array(image_np_or_pil.convert("L"))
            return ocr_with_easyocr(gray)
    else:
        if isinstance(image_np_or_pil, np.ndarray):
            return ocr_with_easyocr(image_np_or_pil)
        else:
            # صورة PIL واردة من الجداول مثلاً، نحولها
            gray = np.array(image_np_or_pil.convert("L"))
            return ocr_with_easyocr(gray)

# ------------------------------
# استخراج الجداول (يستخدم ocr_image للخلايا)
# ------------------------------
def detect_and_extract_tables(image_pil, handwritten=False):
    processor, det_model, struct_model = get_table_models()
    if None in (processor, det_model, struct_model):
        return []
    inputs = processor(images=image_pil, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k,v in inputs.items()}
    with torch.no_grad():
        outputs = det_model(**inputs)
    target_sizes = torch.tensor([image_pil.size[::-1]])
    results = processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
    tables = []
    for box in results["boxes"]:
        xmin, ymin, xmax, ymax = box.tolist()
        cropped = image_pil.crop((xmin, ymin, xmax, ymax))
        table = extract_table_structure(cropped, handwritten)
        if table:
            tables.append(table)
    return tables

def extract_table_structure(cropped_image, handwritten=False):
    processor, det_model, struct_model = get_table_models()
    if None in (processor, det_model, struct_model):
        return None
    inputs = processor(images=cropped_image, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k,v in inputs.items()}
    with torch.no_grad():
        outputs = struct_model(**inputs)
    target_sizes = torch.tensor([cropped_image.size[::-1]])
    results = processor.post_process_object_detection(outputs, threshold=0.85, target_sizes=target_sizes)[0]
    cells = []
    for box, label in zip(results["boxes"], results["labels"]):
        if label in [2,3,4,5,6]:
            xmin, ymin, xmax, ymax = box.tolist()
            cells.append({"xmin":xmin, "ymin":ymin, "xmax":xmax, "ymax":ymax})
    if not cells:
        return None
    cells.sort(key=lambda c: (c["ymin"], c["xmin"]))
    rows, cur_row, last_y = [], [], cells[0]["ymin"]
    for c in cells:
        if abs(c["ymin"] - last_y) > 10:
            cur_row.sort(key=lambda x: x["xmin"])
            rows.append(cur_row)
            cur_row, last_y = [], c["ymin"]
        cur_row.append(c)
    if cur_row:
        cur_row.sort(key=lambda x: x["xmin"])
        rows.append(cur_row)
    # OCR للخلايا (دائماً EasyOCR لأن الجداول ليست خط يد)
    table_matrix = []
    for row in rows:
        row_texts = []
        for cell in row:
            cell_img = cropped_image.crop((cell["xmin"], cell["ymin"], cell["xmax"], cell["ymax"]))
            cell_np = np.array(cell_img.convert("L"))
            cell_gray = preprocess_image(cell_np, heavy=False)
            text = ocr_with_easyocr(cell_gray)
            row_texts.append(text.strip())
        table_matrix.append(row_texts)
    return table_matrix

# ------------------------------
# تحليل بادئات اسم الملف
# ------------------------------
def parse_filename_flags(filename: str):
    base = os.path.splitext(filename)[0]
    if base.startswith("radar_"):
        base = base[6:]
    flags = {
        "force_table": None,
        "handwritten": False,
        "heavy_prep": False,
        "dpi": DEFAULT_DPI,
        "fast": False
    }
    parts = base.split("_")
    i = 0
    while i < len(parts):
        p = parts[i].lower()
        if p == "tbl":
            flags["force_table"] = True
        elif p == "notbl":
            flags["force_table"] = False
        elif p == "hw":
            flags["handwritten"] = True
        elif p == "prep":
            flags["heavy_prep"] = True
        elif p == "dpi300":
            flags["dpi"] = 300
        elif p == "fast":
            flags["fast"] = True
            flags["force_table"] = False
            flags["heavy_prep"] = False
            flags["handwritten"] = False
        else:
            pass
        i += 1
    return flags

# ------------------------------
# المعالجة المركزية (مع دعم الخط اليدوي)
# ------------------------------
def core_ocr_process_and_upload(filename: str, file_path: str):
    global ocr_stats
    token = os.environ.get("AIOCR_KEY")
    if not token:
        logger.error("❌ AIOCR_KEY غير موجود")
        return None
    api = HfApi(token=token)

    flags = parse_filename_flags(filename)
    use_table = ENABLE_TABLE_DEFAULT if flags["force_table"] is None else flags["force_table"]
    use_heavy = flags["heavy_prep"]
    use_handwritten = flags["handwritten"]
    dpi = flags["dpi"]
    fast_mode = flags["fast"]

    clean_name = os.path.splitext(filename)[0].replace("radar_", "")
    final_cloud_filename = f"{clean_name}.txt"
    logger.info(f"🔎 معالجة {clean_name} | جداول:{use_table} يدوي:{use_handwritten} ثقيل:{use_heavy} DPI:{dpi}")

    ocr_stats["status"] = f"🔎 معالجة {clean_name}..."

    final_text = ""
    try:
        if filename.lower().endswith('.pdf'):
            doc = fitz.open(file_path)
            for page_num in range(len(doc)):
                page = doc[page_num]
                if fast_mode and not use_table and not use_handwritten:
                    t = page.get_text()
                    if t and len(t.strip()) > 15:
                        final_text += t + "\n"
                        continue
                pix = page.get_pixmap(dpi=dpi)
                img_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
                if pix.n >= 3:
                    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
                gray = preprocess_image(img_np, heavy=use_heavy)

                # OCR للصفحة
                page_text = ocr_image(gray, handwritten=use_handwritten)
                final_text += page_text + "\n"

                # جداول
                if use_table:
                    try:
                        orig_pil = Image.fromarray(cv2.cvtColor(img_np, cv2.COLOR_GRAY2RGB) if len(img_np.shape)==2 else cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
                        tables = detect_and_extract_tables(orig_pil, handwritten=False)
                        import csv
                        for idx, tbl in enumerate(tables):
                            csv_filename = f"{clean_name}_page{page_num+1}_table{idx+1}.csv"
                            csv_path = f"/tmp/{csv_filename}"
                            with open(csv_path, "w", encoding="utf-8", newline="") as cf:
                                writer = csv.writer(cf)
                                writer.writerows(tbl)
                            api.upload_file(path_or_fileobj=csv_path,
                                            path_in_repo=f"{PROCESS_DIR}/{csv_filename}",
                                            repo_id=REPO_ID, repo_type="dataset", token=token)
                            os.remove(csv_path)
                            logger.info(f"📊 رفع جدول: {csv_filename}")
                    except Exception as e:
                        logger.warning(f"فشل جداول صفحة {page_num+1}: {e}")
                del pix, img_np, gray
                gc.collect()
            doc.close()
        else:
            img = Image.open(file_path)
            img_np = np.array(img.convert("L"))
            gray = preprocess_image(img_np, heavy=use_heavy)

            final_text = ocr_image(gray, handwritten=use_handwritten)

            if use_table:
                try:
                    orig_pil = img.convert("RGB")
                    tables = detect_and_extract_tables(orig_pil, handwritten=False)
                    import csv
                    for idx, tbl in enumerate(tables):
                        csv_filename = f"{clean_name}_table{idx+1}.csv"
                        csv_path = f"/tmp/{csv_filename}"
                        with open(csv_path, "w", encoding="utf-8", newline="") as cf:
                            writer = csv.writer(cf)
                            writer.writerows(tbl)
                        api.upload_file(path_or_fileobj=csv_path,
                                        path_in_repo=f"{PROCESS_DIR}/{csv_filename}",
                                        repo_id=REPO_ID, repo_type="dataset", token=token)
                        os.remove(csv_path)
                        logger.info(f"📊 رفع جدول: {csv_filename}")
                except Exception as e:
                    logger.warning(f"فشل جداول الصورة: {e}")

        if final_text.strip():
            local_txt = f"/tmp/{final_cloud_filename}"
            with open(local_txt, "w", encoding="utf-8") as f:
                f.write(final_text)
            api.upload_file(path_or_fileobj=local_txt,
                            path_in_repo=f"{PROCESS_DIR}/{final_cloud_filename}",
                            repo_id=REPO_ID, repo_type="dataset", token=token)
            os.remove(local_txt)

        ocr_stats["total_extracted_words"] += len(final_text.split())
        ocr_stats["processed_files_count"] += 1
        ocr_stats["last_processed_file"] = f"{clean_name} ({time.strftime('%H:%M:%S')})"
        ocr_stats["micro_logs"] = f"✅ تمت معالجة {clean_name}"
        return final_text
    except Exception as e:
        logger.error(f"❌ فشل {clean_name}: {e}", exc_info=True)
        return None
    finally:
        ocr_stats["status"] = "🟢 المحرك جاهز"
        gc.collect()

# ------------------------------
# الرادار السحابي (بدون تغيير)
# ------------------------------
def cloud_folder_polling_worker():
    while True:
        token = os.environ.get("AIOCR_KEY")
        if token:
            try:
                api = HfApi(token=token)
                all_files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
                valid_files = [f for f in all_files if f.lower().startswith(f"{RAW_DIR}/") and not f.endswith('/')]
                if valid_files:
                    target_file = valid_files[0]
                    filename = os.path.basename(target_file)
                    local_path = os.path.join(UPLOAD_FOLDER, f"radar_{filename}")
                    downloaded = hf_hub_download(repo_id=REPO_ID, filename=target_file, repo_type="dataset", token=token)
                    shutil.copy2(downloaded, local_path)
                    output = core_ocr_process_and_upload(filename, local_path)
                    if output is not None:
                        try:
                            api.delete_file(path_in_repo=target_file, repo_id=REPO_ID, repo_type="dataset", token=token)
                            logger.info(f"🗑️ حذف {filename}")
                        except Exception as e:
                            logger.error(f"⚠️ فشل حذف {filename}: {e}")
                        if os.path.exists(local_path):
                            os.remove(local_path)
            except Exception as e:
                logger.error(f"⚠️ خطأ في الرادار: {e}")
        time.sleep(4)

polling_thread = threading.Thread(target=cloud_folder_polling_worker, daemon=True)
polling_thread.start()

# ------------------------------
# واجهة المراقبة (مصححة وخالية من مشاكل HTML)
# ------------------------------
def get_dashboard_data():
    return (ocr_stats["status"], str(ocr_stats["processed_files_count"]),
            str(ocr_stats["total_extracted_words"]), ocr_stats["last_processed_file"],
            ocr_stats["micro_logs"])

with gr.Blocks(title="Sapphire OCR V9") as demo:
    gr.Markdown("# 🔮 Sapphire OCR V9 (EasyOCR + خط يدوي عربي)")
    gr.Markdown("يدعم الكتابة اليدوية العربية عبر البادئة `hw_`")
    
    status_box = gr.Textbox(label="📡 حالة السيرفر", value=ocr_stats["status"], interactive=False)
    with gr.Row():
        count_box = gr.Textbox(label="✅ ملفات", value="0", interactive=False)
        words_box = gr.Textbox(label="🔍 كلمات", value="0", interactive=False)
        last_file_box = gr.Textbox(label="📁 آخر ملف", value="لا طلبات", interactive=False)
    logs_box = gr.Textbox(label="📝 السجل", value=ocr_stats["micro_logs"], lines=3, interactive=False)
    refresh_btn = gr.Button("🔄 تحديث")
    refresh_btn.click(fn=get_dashboard_data, outputs=[status_box, count_box, words_box, last_file_box, logs_box])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)