Update app.py

app.py

@@ -1,95 +1,177 @@
 import os
 import io
-import tempfile
-import yaml
 import numpy as np
 from PIL import Image
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import HTMLResponse
 from huggingface_hub import hf_hub_download
 
-# === Ограничиваем треды ONNX Runtime ДО импорта rapidocr ===
-# На HF Spaces CPU Basic всего 2 vCPU, лишние треды только мешают
+# Ограничиваем треды ONNX Runtime (на HF Spaces CPU Basic всего 2 vCPU)
 os.environ["OMP_NUM_THREADS"] = "2"
 os.environ["MKL_NUM_THREADS"] = "2"
 
-from rapidocr_onnxruntime import RapidOCR
-
 app = FastAPI()
 MODEL_DIR = os.path.join(os.path.dirname(__file__), "models")
 
+def download_file(repo_id, filename, local_dir):
+    """Скачивает файл с HF Hub, если его нет локально."""
+    local_path = os.path.join(local_dir, filename)
+    if os.path.exists(local_path):
+        return local_path
+    try:
+        print(f"Downloading {repo_id}/{filename}...")
+        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
+        print(f"  Saved to {local_path}")
+        return local_path
+    except Exception as e:
+        print(f"  Failed: {e}")
+        return None
+
 def ensure_models():
-    """Скачивает модели при первом запуске, если их нет локально."""
+    """Скачивает det, rec, dict. Cls скачаем отдельно если нужен config."""
     os.makedirs(MODEL_DIR, exist_ok=True)
-    repo_id = "monkt/paddleocr-onnx"
-    files = [
-        "detection/v5/det.onnx",
-        "languages/eslav/rec.onnx",
-        "languages/eslav/dict.txt",
-    ]
-    for remote_path in files:
-        local_path = os.path.join(MODEL_DIR, remote_path)
-        if not os.path.exists(local_path):
-            print(f"Downloading {remote_path}...")
-            hf_hub_download(repo_id=repo_id, filename=remote_path, local_dir=MODEL_DIR)
-            print(f"  Saved {remote_path}")
-    return {
-        "det": os.path.join(MODEL_DIR, "detection/v5/det.onnx"),
-        "rec": os.path.join(MODEL_DIR, "languages/eslav/rec.onnx"),
-        "dict": os.path.join(MODEL_DIR, "languages/eslav/dict.txt"),
-    }
-
-model_paths = ensure_models()
-
-# === Кастомный config для максимальной скорости на CPU ===
-config = {
-    "Global": {
-        "text_score": 0.5,
-        "use_angle_cls": False,
-        "use_text_det": True,
-        "print_verbose": False,
-        "min_height": 30,
-    },
-    "Det": {
-        "model_path": model_paths["det"],
-        "limit_side_len": 480,      # ← меньше пикселей = быстрее детекция
-        "limit_type": "max",        # ← не растягивает маленькие, сжимает большие
-        "thresh": 0.5,              # ← меньше шумовых областей
-        "box_thresh": 0.5,
-        "max_candidates": 1000,
-        "unclip_ratio": 1.6,
-        "use_dilation": False,      # ← отключено, дилатация тормозит на CPU
-        "score_mode": "fast",
-    },
-    "Rec": {
-        "model_path": model_paths["rec"],
-        "keys_path": model_paths["dict"],
-        "img_shape": [3, 48, 320],
-        "batch_num": 1,             # ← на CPU батчинг не ускоряет, а замедляет
-    },
-    "Cls": {
-        "model_path": "",
-        "img_shape": [3, 48, 192],
-        "label_list": ["0", "180"],
-        "batch_num": 1,
-        "thresh": 0.9,
-    },
-}
-
-with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-    yaml.dump(config, f, default_flow_style=False)
-    config_path = f.name
-
-ocr = RapidOCR(config_path=config_path)
-
-# === Warm-up: разогреваем ONNX Runtime ===
+    det = download_file("monkt/paddleocr-onnx", "detection/v5/det.onnx", MODEL_DIR)
+    rec = download_file("monkt/paddleocr-onnx", "languages/eslav/rec.onnx", MODEL_DIR)
+    dict_path = download_file("monkt/paddleocr-onnx", "languages/eslav/dict.txt", MODEL_DIR)
+    return det, rec, dict_path
+
+det_path, rec_path, dict_path = ensure_models()
+
+from rapidocr_onnxruntime import RapidOCR
+
+# === Попытка 1: использовать kwargs (новая версия 1.3.x) ===
+try:
+    ocr = RapidOCR(
+        det_model_path=det_path,
+        rec_model_path=rec_path,
+        rec_keys_path=dict_path,
+        cls=False,
+    )
+    print("✓ RapidOCR initialized with kwargs")
+
+    # Monkey-patch параметров детекции для ускорения
+    if hasattr(ocr, "text_detector") and hasattr(ocr.text_detector, "preprocess_op"):
+        for op in ocr.text_detector.preprocess_op:
+            if op.__class__.__name__ == "DetResizeForTest":
+                op.limit_side_len = 480
+                op.limit_type = "max"
+                print("  Patched: limit_side_len=480, limit_type=max")
+
+    if hasattr(ocr, "text_detector") and hasattr(ocr.text_detector, "postprocess_op"):
+        ocr.text_detector.postprocess_op.thresh = 0.5
+        ocr.text_detector.postprocess_op.use_dilation = False
+        ocr.text_detector.postprocess_op.score_mode = "fast"
+        print("  Patched: thresh=0.5, use_dilation=False, score_mode=fast")
+
+    if hasattr(ocr, "text_recognizer"):
+        ocr.text_recognizer.rec_batch_num = 1
+        print("  Patched: rec_batch_num=1")
+
+except TypeError:
+    # === Попытка 2: старая версия, нужен полный config.yaml ===
+    print("kwargs not supported, using full config.yaml")
+
+    # Скачиваем cls модель из нескольких возможных источников
+    cls_path = None
+    for repo, path in [
+        ("RapidAI/RapidOCR", "onnx/PP-OCRv4/cls/ch_ppocr_mobile_v2.0_cls_infer.onnx"),
+        ("RapidAI/RapidOCR", "onnx/PP-OCRv3/cls/ch_ppocr_mobile_v2.0_cls_infer.onnx"),
+        ("RapidAI/RapidOCR", "resources/models/ch_ppocr_mobile_v2.0_cls_infer.onnx"),
+    ]:
+        cls_path = download_file(repo, path, MODEL_DIR)
+        if cls_path:
+            break
+
+    # Если не нашли — запускаем стандартный RapidOCR, чтобы он скачал модели в кэш
+    if not cls_path:
+        print("Downloading standard models via RapidOCR()...")
+        temp_ocr = RapidOCR()
+        import rapidocr_onnxruntime
+        pkg_dir = os.path.dirname(rapidocr_onnxruntime.__file__)
+        for root, dirs, files in os.walk(pkg_dir):
+            for f in files:
+                if "cls" in f and f.endswith(".onnx"):
+                    cls_path = os.path.join(root, f)
+                    break
+            if cls_path:
+                break
+        # Ищем в кэше
+        if not cls_path:
+            cache_dir = os.path.expanduser("~/.cache/rapidocr_onnxruntime")
+            if os.path.exists(cache_dir):
+                for root, dirs, files in os.walk(cache_dir):
+                    for f in files:
+                        if "cls" in f and f.endswith(".onnx"):
+                            cls_path = os.path.join(root, f)
+                            break
+                    if cls_path:
+                        break
+
+    # Создаём полный config.yaml со ВСЕМИ обязательными полями
+    config_path = os.path.join(os.path.dirname(__file__), "config.yaml")
+    with open(config_path, "w") as f:
+        f.write(f"""
+Global:
+  text_score: 0.5
+  use_angle_cls: false
+  print_verbose: false
+  min_height: 30
+  width_height_ratio: 8
+
+Det:
+  module_name: ch_ppocr_v3_det
+  class_name: TextDetector
+  model_path: {det_path}
+  use_cuda: false
+  pre_process:
+    DetResizeForTest:
+      limit_side_len: 480
+      limit_type: max
+    NormalizeImage:
+      std: [0.229, 0.224, 0.225]
+      mean: [0.485, 0.456, 0.406]
+      scale: 1./255.
+      order: hwc
+    ToCHWImage:
+    KeepKeys:
+      keep_keys: ['image', 'shape']
+  post_process:
+    thresh: 0.5
+    box_thresh: 0.5
+    max_candidates: 1000
+    unclip_ratio: 1.6
+    use_dilation: false
+    score_mode: fast
+
+Cls:
+  module_name: ch_ppocr_v2_cls
+  class_name: TextClassifier
+  model_path: {cls_path or det_path}
+  cls_img_shape: [3, 48, 192]
+  cls_batch_num: 1
+  cls_thresh: 0.9
+  label_list: [0, 180]
+
+Rec:
+  module_name: ch_ppocr_v2_rec
+  class_name: TextRecognizer
+  model_path: {rec_path}
+  rec_img_shape: [3, 48, 320]
+  rec_batch_num: 1
+  keys_path: {dict_path}
+""")
+
+    ocr = RapidOCR(config_path=config_path)
+    print("✓ RapidOCR initialized with config.yaml")
+
+# Warm-up: разогреваем ONNX Runtime
 print("Warming up OCR...")
 warmup_img = np.zeros((480, 480, 3), dtype=np.uint8)
 try:
     _ = ocr(warmup_img)
 except Exception:
     pass
-print("Warm-up done!")
+print("Ready!")
 
 HTML_FORM = """<!DOCTYPE html>
 <html>
@@ -141,20 +223,14 @@ async def predict(file: UploadFile = File(...)):
     try:
         contents = await file.read()
         image = Image.open(io.BytesIO(contents))
-        
         if image.mode != "RGB":
             image = image.convert("RGB")
-        
-        # Оптимизация: уменьшаем до 480px по длинной стороне
         if image.width > 480 or image.height > 480:
             image.thumbnail((480, 480))
-            
         img_array = np.array(image)
         results, _ = ocr(img_array)
-        
         if not results:
             return {"text": ""}
-            
         text = " ".join([item[1] for item in results]).strip()
         return {"text": text}
     except Exception as e: