edit

app.py

@@ -7,13 +7,12 @@ import soundfile as soundfile_module
 import torch
 import gradio as gradio_module
 from PIL import Image
+import easyocr
 from transformers import (
     pipeline,
     VitsModel,
     AutoTokenizer,
 )
-from nemotron_ocr.inference.pipeline import NemotronOCR  # <-- Nemotron OCR v1
-
 
 # ============================
 # 1. Настройки устройства
@@ -23,76 +22,87 @@ device_string: str = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 # ============================
-# 2. Модели
+# 2. OCR (easyocr, английский)
 # ============================
 
-ocr_engine: NemotronOCR = NemotronOCR()
-
-summary_pipeline = pipeline(
-    task="summarization",
-    model="sshleifer/distilbart-cnn-12-6",
+# TODO_USER: при желании можно добавить другие языки, но тогда конспект и TTS всё равно останутся на английском
+ocr_reader = easyocr.Reader(
+    ["en"],  # языки
+    gpu=(device_string == "cuda"),
 )
 
-tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-eng")
-tts_tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
-tts_model.to(device_string)
-
-
-# ============================
-# 3. OCR через NemotronOCR
-# ============================
 
 def run_ocr(image_object: Image.Image) -> str:
     """
-    OCR для печатного (и вообще любого) английского текста с картины.
-
-    Используем NemotronOCR из nvidia/nemotron-ocr-v1.
-    Модель сама делает:
-    - детекцию текстовых блоков,
-    - распознавание текста,
-    - анализ порядка чтения.
-
-    На выходе NemotronOCR даёт список dict:
-    [
-        {
-            "text": "...",
-            "confidence": float,
-            "left": float,
-            "upper": float,
-            "right": float,
-            "lower": float,
-            ...
-        },
-        ...
-    ]
+    OCR для печатного английского текста.
+    Используем easyocr, потому что он реально более устойчивый для
+    произвольных сканов/фото, чем большинство трансформеров, которые мы пробовали.
     """
-
     if image_object is None:
         return ""
 
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temporary_file:
-        image_object.save(temporary_file.name)
-        image_path: str = temporary_file.name
+    rgb_image_object: Image.Image = image_object.convert("RGB")
 
-    predictions = ocr_engine(image_path)
+    # easyocr работает с numpy-массивом
+    numpy_image = numpy_module.array(rgb_image_object)
+
+    results = ocr_reader.readtext(
+        numpy_image,
+        detail=1,   # возвращаем bbox + текст + confidence
+        paragraph=True,  # склеивать текст в параграфы, где это возможно
+    )
 
     text_parts = []
-    for prediction in predictions:
-        text_value = prediction.get("text", "")
+    for bbox, text_value, confidence_value in results:
         if not text_value:
             continue
-
-
-        text_parts.append(str(text_value))
+        # TODO_USER: при желании можно фильтровать по confidence_value
+        text_parts.append(text_value)
 
     recognized_text: str = "\n".join(text_parts).strip()
     return recognized_text
 
 
 # ============================
-# 4. Суммаризация (английский)
+# 3. Трансформер #1: классификация текста
 # ============================
 
+text_classifier_pipeline = pipeline(
+    task="text-classification",
+    model="distilbert-base-uncased-finetuned-sst-2-english",
+)
+
+
+def run_text_classification(input_text: str) -> str:
+    """
+    Пример анализа текста трансформером:
+    используем sentiment-классификатор как демонстрацию.
+    Возвращаем строку вида: "label: POSITIVE, score: 0.98".
+    """
+    cleaned_text: str = input_text.strip()
+    if not cleaned_text:
+        return ""
+
+    result_list = text_classifier_pipeline(cleaned_text)
+    result = result_list[0]
+
+    label_value: str = str(result.get("label", ""))
+    score_value: float = float(result.get("score", 0.0))
+
+    classification_text: str = f"{label_value} (score={score_value:.3f})"
+    return classification_text
+
+
+# ============================
+# 4. Трансформер #2: суммаризация (английский)
+# ============================
+
+summary_pipeline = pipeline(
+    task="summarization",
+    model="sshleifer/distilbart-cnn-12-6",
+)
+
+
 def run_summarization(
     input_text: str,
     max_summary_tokens: int = 128,
@@ -106,13 +116,13 @@ def run_summarization(
         return ""
 
     word_count: int = len(cleaned_text.split())
-
     dynamic_max_length: int = min(
         max_summary_tokens,
         max(32, word_count + 20),
     )
 
     if word_count < 8:
+        # TODO_USER: для очень короткого текста суммаризация сомнительна, возвращаем исходный текст
         return cleaned_text
 
     summary_result_list = summary_pipeline(
@@ -127,9 +137,14 @@ def run_summarization(
 
 
 # ============================
-# 5. TTS (английский, MMS VITS)
+# 5. Трансформер #3: TTS (английский, MMS VITS)
 # ============================
 
+tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-eng")
+tts_tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+tts_model.to(device_string)
+
+
 def run_tts(summary_text: str) -> Optional[str]:
     """
     Озвучка английского текста конспекта через VitsModel (facebook/mms-tts-eng).
@@ -151,9 +166,7 @@ def run_tts(summary_text: str) -> Optional[str]:
     }
 
     input_ids_tensor = tokenized_inputs.get("input_ids")
-    if input_ids_tensor is None:
-        return None
-    if input_ids_tensor.numel() == 0 or input_ids_tensor.shape[1] == 0:
+    if input_ids_tensor is None or input_ids_tensor.numel() == 0:
         return None
 
     try:
@@ -188,15 +201,18 @@ def run_tts(summary_text: str) -> Optional[str]:
 def full_flow(
     image_object: Image.Image,
     max_summary_tokens: int = 128,
-) -> Tuple[str, str, Optional[str]]:
+) -> Tuple[str, str, str, Optional[str]]:
     """
     Полный пайплайн:
-    1) OCR: изображение -> исходный английский текст
-    2) Суммаризация: текст -> конспект (английский)
-    3) TTS: конспект -> .wav файл (или None, если TTS не смог)
+    1) OCR (easyocr): изображение -> исходный текст (английский)
+    2) Классификация текста трансформером (sentiment)
+    3) Суммаризация: текст -> конспект
+    4) TTS: конспект -> .wav файл (или None)
     """
     recognized_text: str = run_ocr(image_object=image_object)
 
+    classification_text: str = run_text_classification(recognized_text)
+
     summary_text: str = run_summarization(
         input_text=recognized_text,
         max_summary_tokens=max_summary_tokens,
@@ -204,7 +220,7 @@ def full_flow(
 
     audio_file_path: Optional[str] = run_tts(summary_text=summary_text)
 
-    return recognized_text, summary_text, audio_file_path
+    return recognized_text, classification_text, summary_text, audio_file_path
 
 
 # ============================
@@ -228,25 +244,29 @@ gradio_interface = gradio_module.Interface(
     ],
     outputs=[
         gradio_module.Textbox(
-            label="Распознанный текст (Nemotron OCR)",
+            label="Распознанный текст (OCR, easyocr)",
             lines=8,
         ),
         gradio_module.Textbox(
-            label="Конспект (английский текст)",
+            label="Анализ текста (классификация, DistilBERT)",
+            lines=2,
+        ),
+        gradio_module.Textbox(
+            label="Конспект (английский текст, DistilBART)",
             lines=6,
         ),
         gradio_module.Audio(
-            label="Озвучка конспекта (английский TTS)",
+            label="Озвучка конспекта (английский TTS, VITS)",
             type="filepath",
         ),
     ],
-    title="Картинка → Текст → Конспект → Озвучка (Nemotron OCR + английские модели)",
+    title="Картинка → Текст → Анализ → Конспект → Озвучка",
     description=(
-        "1) Nemotron OCR v1 (nvidia/nemotron-ocr-v1) распознаёт текст с документа.\n"
-        "2) Английский трансформер суммаризации делает краткий пересказ.\n"
-        "3) VITS-модель MMS (facebook/mms-tts-eng) озвучивает конспект.\n\n"
-        "Если озвучка не сгенерировалась, значит конкретный текст не понравился TTS-модели "
-        "и она упала внутри — пайплайн просто пропустит аудио."
+        "1) easyocr распознаёт печатный английский текст с картинки.\n"
+        "2) Трансформер-классификатор (DistilBERT) оценивает тон текста.\n"
+        "3) Трансформер-суммаризатор (DistilBART) делает краткий конспект.\n"
+        "4) Трансформер TTS (MMS VITS) озвучивает конспект.\n"
+        "В проекте используются три трансф��рмера с Hugging Face, OCR сделан через easyocr."
     ),
 )
 

requirements.txt

@@ -1,7 +1,8 @@
-transformers>=4.40.0
+transformers>=4.33.0
 torch
 sentencepiece
 gradio
 Pillow
 numpy
 soundfile
+easyocr