update

app.py

@@ -2,6 +2,7 @@ import os
 import gc
 import json
 import time
+import traceback
 from pathlib import Path
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple
@@ -23,8 +24,8 @@ DET_FALLBACK_PT = os.getenv("DET_FALLBACK_PT", "best.pt")
 
 CLS_REPO_ID = os.getenv("CLS_REPO_ID", "mr3vial/paleo-hebrew-convnext")
 CLS_WEIGHTS = os.getenv("CLS_WEIGHTS", "convnext_weights.pt") 
-CLS_CLASSES = os.getenv("CLS_CLASSES", "classes.json")
-CLS_MODEL_NAME = os.getenv("CLS_MODEL_NAME", "convnext_large")
+CLS_CLASSES = os.getenv("CLS_CLASSES", "classes.txt")
+CLS_MODEL_NAME = os.getenv("CLS_MODEL_NAME", "convnext_base")
 
 MT5_BOX2HE_REPO = os.getenv("MT5_BOX2HE_REPO", "mr3vial/paleo-hebrew-mt5-post-ocr-processing")
 MT5_BOX2EN_REPO = os.getenv("MT5_BOX2EN_REPO", "mr3vial/paleo-hebrew-mt5-translate")
@@ -51,10 +52,6 @@ def safe_int(x: float, default: int = 0) -> int:
         return default
 
 def sort_boxes_reading_order(boxes_xyxy: List[List[float]], rtl: bool = True) -> List[int]:
-    """
-    Rough reading order:
-      - sort by y-center, then within a line by x-center (rtl: desc, ltr: asc).
-    """
     if not boxes_xyxy:
         return []
     centers = []
@@ -102,12 +99,22 @@ def load_detector():
 @torch.no_grad()
 def load_classifier():
     import timm
-    from safetensors.torch import load_file as safetensors_load
+    
     device = get_device()
 
     classes_path = hf_hub_download(repo_id=CLS_REPO_ID, filename=CLS_CLASSES)
-    with open(classes_path, "r", encoding="utf-8") as f:
-        raw = [ln.strip() for ln in f if ln.strip()]
+    
+    # Безопасное чтение классов (поддержка и .json, и .txt)
+    if classes_path.endswith(".json"):
+        with open(classes_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            if isinstance(data, dict):
+                raw = [str(v) for v in data.values()]
+            else:
+                raw = [str(x) for x in data]
+    else:
+        with open(classes_path, "r", encoding="utf-8") as f:
+            raw = [ln.strip() for ln in f if ln.strip()]
 
     letters: List[str] = []
     for ln in raw:
@@ -118,7 +125,10 @@ def load_classifier():
 
     weights_path = hf_hub_download(repo_id=CLS_REPO_ID, filename=CLS_WEIGHTS)
     model = timm.create_model(CLS_MODEL_NAME, pretrained=False, num_classes=len(letters))
-    state = safetensors_load(weights_path)
+    
+    # Загрузка классических .pt весов вместо safetensors
+    ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
+    state = ckpt.get("model", ckpt.get("state_dict", ckpt))
 
     new_state = {}
     for k, v in state.items():
@@ -141,7 +151,6 @@ def preprocess_crop_imagenet(pil: Image.Image) -> torch.Tensor:
 @torch.no_grad()
 def load_mt5(repo_id: str):
     device = get_device()
-    # Memory optimization: Load in bfloat16 to fit multiple models in Space RAM
     dtype = torch.bfloat16 if (device == "cpu" or torch.cuda.is_bf16_supported()) else torch.float16
     tok = AutoTokenizer.from_pretrained(repo_id, use_fast=False)
     model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -157,251 +166,89 @@ def mt5_generate(tok, model, text: str, max_new_tokens: int = 128) -> str:
     out = model.generate(**inp, max_new_tokens=max_new_tokens, num_beams=4, do_sample=False)
     return normalize_spaces(tok.batch_decode(out, skip_special_tokens=True)[0])
 
-# =============================================================================
-# Optional Hebrew->English translators (user-selectable)
-# =============================================================================
-TRANSLATOR_CHOICES = [
-    ("None (use direct box→en)", "none"),
-    ("OPUS-MT he→en (Helsinki-NLP/opus-mt-tc-big-he-en)", "opus"),
-    ("NLLB-200 distilled 600M (he→en) [CC-BY-NC]", "nllb"),
-    ("M2M100 418M (he→en)", "m2m100"),
-]
-
-_TRANSLATOR_CACHE: Dict[str, Tuple[Any, Any, str]] = {}  # kind -> (tok, model, kind)
-
-@torch.no_grad()
-def get_he2en_translator(kind: str):
-    global _TRANSLATOR_CACHE
-    if kind in _TRANSLATOR_CACHE:
-        return _TRANSLATOR_CACHE[kind]
-
-    # Memory Safety: Clear old translators from memory before loading a new one
-    _TRANSLATOR_CACHE.clear()
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    device = get_device()
-
-    if kind == "opus":
-        from transformers import MarianMTModel, MarianTokenizer
-        repo = "Helsinki-NLP/opus-mt-tc-big-he-en"
-        tok = MarianTokenizer.from_pretrained(repo)
-        model = MarianMTModel.from_pretrained(repo).to(device).eval()
-
-    elif kind == "nllb":
-        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-        repo = "facebook/nllb-200-distilled-600M"
-        tok = AutoTokenizer.from_pretrained(repo, src_lang="heb_Hebr")
-        model = AutoModelForSeq2SeqLM.from_pretrained(repo).to(device).eval()
-
-    elif kind == "m2m100":
-        from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
-        repo = "facebook/m2m100_418M"
-        tok = M2M100Tokenizer.from_pretrained(repo)
-        tok.src_lang = "he"
-        model = M2M100ForConditionalGeneration.from_pretrained(repo).to(device).eval()
-
-    else:
-        _TRANSLATOR_CACHE[kind] = (None, None, "none")
-        return _TRANSLATOR_CACHE[kind]
-
-    _TRANSLATOR_CACHE[kind] = (tok, model, kind)
-    return _TRANSLATOR_CACHE[kind]
-
-@torch.no_grad()
-def translate_he_to_en(text_he: str, kind: str) -> str:
-    if kind == "none":
-        return ""
-    tok, model, kind = get_he2en_translator(kind)
-    device = get_device()
-
-    if kind == "opus":
-        batch = tok([text_he], return_tensors="pt", padding=True, truncation=True).to(device)
-        out = model.generate(**batch, max_new_tokens=128, num_beams=4)
-        return normalize_spaces(tok.batch_decode(out, skip_special_tokens=True)[0])
-
-    if kind == "nllb":
-        batch = tok([text_he], return_tensors="pt", truncation=True, max_length=512).to(device)
-        out = model.generate(**batch, forced_bos_token_id=tok.lang_code_to_id["eng_Latn"], max_new_tokens=128, num_beams=4)
-        return normalize_spaces(tok.batch_decode(out, skip_special_tokens=True)[0])
-
-    if kind == "m2m100":
-        batch = tok([text_he], return_tensors="pt", truncation=True, max_length=512).to(device)
-        out = model.generate(**batch, forced_bos_token_id=tok.get_lang_id("en"), max_new_tokens=128, num_beams=4)
-        return normalize_spaces(tok.batch_decode(out, skip_special_tokens=True)[0])
-
-    return ""
-
-# =============================================================================
-# Pipeline Components
-# =============================================================================
-@dataclass
-class Loaded:
-    det: Any
-    det_name: str
-    cls: Any
-    cls_letters: List[str]
-    cls_name: str
-    mt5_he_tok: Any
-    mt5_he: Any
-    mt5_en_tok: Any
-    mt5_en: Any
-
-LOADED: Optional[Loaded] = None
-
-def ensure_loaded() -> Loaded:
-    global LOADED
-    if LOADED is not None:
-        return LOADED
-
-    det, det_name = load_detector()
-    cls, letters, cls_name = load_classifier()
-    he_tok, he_model = load_mt5(MT5_BOX2HE_REPO)
-    en_tok, en_model = load_mt5(MT5_BOX2EN_REPO)
-
-    LOADED = Loaded(
-        det=det, det_name=det_name,
-        cls=cls, cls_letters=letters, cls_name=cls_name,
-        mt5_he_tok=he_tok, mt5_he=he_model,
-        mt5_en_tok=en_tok, mt5_en=en_model,
-    )
-    return LOADED
-
-def yolo_predict(det_model, pil: Image.Image, conf: float, iou: float, max_det: int):
-    # Direct inference, no tempfile needed
-    res = det_model.predict(
-        source=pil,
-        conf=float(conf),
-        iou=float(iou),
-        max_det=int(max_det),
-        verbose=False,
-    )
-    r0 = res[0]
-    if r0.boxes is None or len(r0.boxes) == 0:
-        return [], []
-    xyxy = r0.boxes.xyxy.detach().float().cpu().numpy().tolist()
-    scores = r0.boxes.conf.detach().float().cpu().numpy().tolist()
-    return xyxy, scores
-
-@torch.no_grad()
-def classify_crops(cls_model, letters: List[str], pil: Image.Image, boxes_xyxy, pad: float, topk: int):
-    device = get_device()
-    W, H = pil.size
-    crops = []
-    for (x1, y1, x2, y2) in boxes_xyxy:
-        w = max(1.0, x2 - x1)
-        h = max(1.0, y2 - y1)
-        p = int(round(max(w, h) * float(pad)))
-        a, b = max(0, safe_int(x1) - p), max(0, safe_int(y1) - p)
-        c, d = min(W, safe_int(x2) + p), min(H, safe_int(y2) + p)
-        crops.append(pil.crop((a, b, c, d)).resize((224, 224), Image.BICUBIC))
-
-    if not crops:
-        return [], [], []
-
-    batch = torch.stack([preprocess_crop_imagenet(c) for c in crops], dim=0).to(device)
-    logits = cls_model(batch)
-    probs = torch.softmax(logits, dim=-1).detach().float().cpu().numpy()
-
-    topk_list, top1_letters = [], []
-    for i in range(len(crops)):
-        p = probs[i]
-        idx = np.argsort(-p)[: max(1, int(topk))]
-        row = [(letters[j], float(p[j])) for j in idx]
-        topk_list.append(row)
-        top1_letters.append(row[0][0] if row else "?")
-
-    crop_gallery = [(crops[i], f"{i+1:02d}: {top1_letters[i]}") for i in range(len(crops))]
-    return topk_list, top1_letters, crop_gallery
-
-def build_source_text(boxes_xyxy: List[List[float]], scores: List[float], topk: List[List[Tuple[str, float]]], rtl: bool) -> str:
-    idxs = sort_boxes_reading_order(boxes_xyxy, rtl=rtl)
-    lines = ["[BOXES_CLS_V3]", f"n={len(idxs)} coord_norm=none"]
-    for j, i in enumerate(idxs, start=1):
-        x1, y1, x2, y2 = boxes_xyxy[i]
-        sc = scores[i] if i < len(scores) else 1.0
-        cands = topk[i] if i < len(topk) else []
-        cls_str = " ".join([f"{ch}:{p:.3f}" for ch, p in cands]) if cands else "?"
-        lines.append(f"{j:02d} x1={x1:.1f} y1={y1:.1f} x2={x2:.1f} y2={y2:.1f} score={sc:.3f} | cls {cls_str}")
-    return "\n".join(lines)
-
-def make_annotated(pil: Image.Image, boxes_xyxy, top1_letters, scores):
-    img = pil.copy()
-    draw = ImageDraw.Draw(img)
-    for i, bb in enumerate(boxes_xyxy):
-        x1, y1, x2, y2 = [safe_int(v) for v in bb]
-        lab = top1_letters[i] if i < len(top1_letters) else "?"
-        sc = scores[i] if i < len(scores) else 0.0
-        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-        text_label = f"{i+1:02d}" + (f" {lab}" if lab != "?" else "")
-        draw.text((x1, y1), text_label, fill="red")
-    return img
-
 # =============================================================================
 # Tab Action Handlers
 # =============================================================================
 def run_pipeline_tab(pil, conf, iou, max_det, crop_pad, topk_k, rtl, output_mode, he2en_kind):
-    if pil is None: return None, [], "", ""
-    L = ensure_loaded()
-
-    boxes, det_scores = yolo_predict(L.det, pil, conf, iou, max_det)
-    if not boxes: return pil, [], "No boxes detected.", "No boxes detected."
-
-    topk_list, top1_letters, crop_gallery = classify_crops(L.cls, L.cls_letters, pil, boxes, pad=crop_pad, topk=topk_k)
-    src = build_source_text(boxes, det_scores, topk_list, rtl=rtl)
-
-    heb = ""
-    eng = ""
-    if output_mode == "he":
-        heb = mt5_generate(L.mt5_he_tok, L.mt5_he, src)
-    elif output_mode == "en_direct":
-        eng = mt5_generate(L.mt5_en_tok, L.mt5_en, src)
-    else: # he_then_en
-        heb = mt5_generate(L.mt5_he_tok, L.mt5_he, src)
-        if he2en_kind != "none":
-            eng = translate_he_to_en(heb, he2en_kind)
-        else:
+    try:
+        if pil is None: return None, [], "", ""
+        pil = pil.convert("RGB") # Защита от RGBA картинок
+        
+        L = ensure_loaded()
+
+        boxes, det_scores = yolo_predict(L.det, pil, conf, iou, max_det)
+        if not boxes: return pil, [], "No boxes detected.", "No boxes detected."
+
+        topk_list, top1_letters, crop_gallery = classify_crops(L.cls, L.cls_letters, pil, boxes, pad=crop_pad, topk=topk_k)
+        src = build_source_text(boxes, det_scores, topk_list, rtl=rtl)
+
+        heb = ""
+        eng = ""
+        if output_mode == "he":
+            heb = mt5_generate(L.mt5_he_tok, L.mt5_he, src)
+        elif output_mode == "en_direct":
             eng = mt5_generate(L.mt5_en_tok, L.mt5_en, src)
-
-    annotated = make_annotated(pil, boxes, top1_letters, det_scores)
-    return annotated, crop_gallery, heb, eng
+        else: # he_then_en
+            heb = mt5_generate(L.mt5_he_tok, L.mt5_he, src)
+            if he2en_kind != "none":
+                eng = translate_he_to_en(heb, he2en_kind)
+            else:
+                eng = mt5_generate(L.mt5_en_tok, L.mt5_en, src)
+
+        annotated = make_annotated(pil, boxes, top1_letters, det_scores)
+        return annotated, crop_gallery, heb, eng
+    except Exception as e:
+        err_msg = f"ERROR:\n{traceback.format_exc()}"
+        return None, [], err_msg, err_msg
 
 def run_detector_tab(pil, conf, iou, max_det):
-    if pil is None: return None, "{}"
-    L = ensure_loaded()
-    boxes, scores = yolo_predict(L.det, pil, conf, iou, max_det)
-    annotated = make_annotated(pil, boxes, [], scores)
-    
-    debug = {"num_boxes": len(boxes), "boxes_xyxy": boxes, "confidences": scores}
-    return annotated, json.dumps(debug, indent=2)
+    try:
+        if pil is None: return None, "{}"
+        pil = pil.convert("RGB")
+        
+        L = ensure_loaded()
+        boxes, scores = yolo_predict(L.det, pil, conf, iou, max_det)
+        annotated = make_annotated(pil, boxes, [], scores)
+        
+        debug = {"num_boxes": len(boxes), "boxes_xyxy": boxes, "confidences": scores}
+        return annotated, json.dumps(debug, indent=2)
+    except Exception as e:
+        return None, json.dumps({"error": traceback.format_exc()}, indent=2)
 
 def run_classifier_tab(pil, conf, iou, max_det, crop_pad, topk_k):
-    if pil is None: return [], "{}"
-    L = ensure_loaded()
-    boxes, _ = yolo_predict(L.det, pil, conf, iou, max_det)
-    if not boxes: return [], "No boxes found to classify."
-    
-    topk_list, top1_letters, crop_gallery = classify_crops(L.cls, L.cls_letters, pil, boxes, pad=crop_pad, topk=topk_k)
-    
-    details = {}
-    for i, (row, box) in enumerate(zip(topk_list, boxes)):
-        details[f"Box_{i+1:02d}"] = {
-            "top_predictions": dict(row),
-            "coordinates": box
-        }
-    return crop_gallery, json.dumps(details, ensure_ascii=False, indent=2)
+    try:
+        if pil is None: return [], "{}"
+        pil = pil.convert("RGB")
+        
+        L = ensure_loaded()
+        boxes, _ = yolo_predict(L.det, pil, conf, iou, max_det)
+        if not boxes: return [], "No boxes found to classify."
+        
+        topk_list, top1_letters, crop_gallery = classify_crops(L.cls, L.cls_letters, pil, boxes, pad=crop_pad, topk=topk_k)
+        
+        details = {}
+        for i, (row, box) in enumerate(zip(topk_list, boxes)):
+            details[f"Box_{i+1:02d}"] = {
+                "top_predictions": dict(row),
+                "coordinates": box
+            }
+        return crop_gallery, json.dumps(details, ensure_ascii=False, indent=2)
+    except Exception as e:
+        return [], json.dumps({"error": traceback.format_exc()}, indent=2)
 
 def save_feedback(heb_pred, eng_pred, heb_corr, eng_corr, notes):
-    rec = {
-        "ts": now_ts(),
-        "heb_pred": heb_pred, "eng_pred": eng_pred,
-        "heb_corr": heb_corr, "eng_corr": eng_corr,
-        "notes": notes,
-    }
-    with open(FEEDBACK_PATH, "a", encoding="utf-8") as f:
-        f.write(json.dumps(rec, ensure_ascii=False) + "\n")
-    return f"Saved to {FEEDBACK_PATH}"
+    try:
+        rec = {
+            "ts": now_ts(),
+            "heb_pred": heb_pred, "eng_pred": eng_pred,
+            "heb_corr": heb_corr, "eng_corr": eng_corr,
+            "notes": notes,
+        }
+        with open(FEEDBACK_PATH, "a", encoding="utf-8") as f:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        return f"Saved to {FEEDBACK_PATH}"
+    except Exception as e:
+        return f"Error saving: {str(e)}"
 
 # =============================================================================
 # UI
@@ -418,7 +265,6 @@ with gr.Blocks(title="Paleo-Hebrew Tablet Reader") as demo:
     gr.Markdown("# Paleo-Hebrew Epigraphy Pipeline")
     
     with gr.Row():
-        # LEFT COLUMN: Global Inputs & Settings
         with gr.Column(scale=3):
             inp = gr.Image(type="pil", label="Input Image")
             
@@ -438,11 +284,8 @@ with gr.Blocks(title="Paleo-Hebrew Tablet Reader") as demo:
             if ex:
                 gr.Examples(examples=ex, inputs=[inp], cache_examples=False)
 
-        # RIGHT COLUMN: Tabs for different pipeline stages
         with gr.Column(scale=7):
             with gr.Tabs():
-                
-                # TAB 1: Main E2E Pipeline
                 with gr.Tab("End-to-End Pipeline"):
                     run_pipe_btn = gr.Button("Run Full Pipeline", variant="primary")
                     out_annot_pipe = gr.Image(type="pil", label="Detections (bbox + top1)")
@@ -451,19 +294,16 @@ with gr.Blocks(title="Paleo-Hebrew Tablet Reader") as demo:
                         out_en_pipe = gr.Textbox(label="English Output", lines=3, interactive=False)
                     out_crops_pipe = gr.Gallery(label="Letter crops").style(grid=6, height="auto") if hasattr(gr.Gallery, 'style') else gr.Gallery(label="Letter crops")
 
-                # TAB 2: Detector Only
                 with gr.Tab("Detector (YOLO)"):
                     run_det_btn = gr.Button("Run Detector", variant="primary")
                     out_annot_det = gr.Image(type="pil", label="Detected Bounding Boxes")
-                    out_json_det = gr.Code(label="Raw Detection Output", language="json")
+                    out_json_det = gr.Code(label="Raw Detection Output / Error Logs", language="json")
 
-                # TAB 3: Classifier Only
                 with gr.Tab("Classifier (ConvNeXt)"):
                     run_cls_btn = gr.Button("Run Classifier", variant="primary")
                     out_crops_cls = gr.Gallery(label="Isolated Letter Crops").style(grid=6, height="auto") if hasattr(gr.Gallery, 'style') else gr.Gallery(label="Isolated Letter Crops")
-                    out_json_cls = gr.Code(label="Top-K Probabilities per Box", language="json")
+                    out_json_cls = gr.Code(label="Top-K Probabilities per Box / Error Logs", language="json")
 
-                # TAB 4: Feedback
                 with gr.Tab("Feedback"):
                     gr.Markdown("Submit corrections to improve the dataset.")
                     heb_corr = gr.Textbox(label="Correct Hebrew", lines=2)
@@ -472,7 +312,6 @@ with gr.Blocks(title="Paleo-Hebrew Tablet Reader") as demo:
                     save_btn = gr.Button("Submit Feedback")
                     save_status = gr.Textbox(label="Status", interactive=False)
 
-    # --- Button Clicks ---
     run_pipe_btn.click(
         fn=run_pipeline_tab,
         inputs=[inp, conf, iou, max_det, crop_pad, topk_k, rtl, output_mode, he2en_kind],