edit

app.py

@@ -7,7 +7,6 @@ import soundfile as soundfile_module
 import torch
 import gradio as gradio_module
 from PIL import Image
-import easyocr
 from transformers import (
     pipeline,
     VitsModel,
@@ -18,83 +17,61 @@ from transformers import (
 # 1. Настройки устройства
 # ============================
 
-device_string: str = "cuda" if torch.cuda.is_available() else "cpu"
+# TODO_USER: для нормальной работы olmOCR почти наверняка нужен GPU
+if torch.cuda.is_available():
+    device_string: str = "cuda"
+    pipeline_device_index: int = 0
+else:
+    device_string = "cpu"
+    pipeline_device_index = -1  # Gradio/transformers: -1 = CPU
 
 
 # ============================
-# 2. OCR (easyocr, английский)
+# 2. OCR на olmOCR-2-7B-1025-FP8
 # ============================
 
-# TODO_USER: при желании можно добавить другие языки, но тогда конспект и TTS всё равно останутся на английском
-ocr_reader = easyocr.Reader(
-    ["en"],  # языки
-    gpu=(device_string == "cuda"),
+# Модель: allenai/olmOCR-2-7B-1025-FP8
+# По README это image-to-text трансформер, так что используем стандартный pipeline.
+ocr_pipeline = pipeline(
+    task="image-to-text",
+    model="allenai/olmOCR-2-7B-1025-FP8",
+    device=pipeline_device_index,
+    # TODO_USER: при необходимости можно добавить torch_dtype=..., но лучше сначала проверить дефолт
 )
 
 
 def run_ocr(image_object: Image.Image) -> str:
     """
-    OCR для печатного английского текста.
-    Используем easyocr, потому что он реально более устойчивый для
-    произвольных сканов/фото, чем большинство трансформеров, которые мы пробовали.
+    OCR для печатного английского текста с помощью olmOCR-2-7B-1025-FP8.
+
+    Вход: PIL.Image (страница/скриншот).
+    Выход: строка текста, которую модель сгенерировала как распознавание.
     """
     if image_object is None:
         return ""
 
     rgb_image_object: Image.Image = image_object.convert("RGB")
 
-    # 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 bbox, text_value, confidence_value in results:
-        if not text_value:
-            continue
-        # TODO_USER: при желании можно фильтровать по confidence_value
-        text_parts.append(text_value)
-
-    recognized_text: str = "\n".join(text_parts).strip()
+    # olmOCR поддерживает прямой вызов через pipeline("image-to-text").
+    # Ожидаемый формат ответа: список dict вида [{"generated_text": "..."}].
+    result = ocr_pipeline(rgb_image_object)
+
+    if isinstance(result, list) and len(result) > 0:
+        first_item = result[0]
+        if isinstance(first_item, dict) and "generated_text" in first_item:
+            text_value: str = str(first_item["generated_text"])
+        else:
+            # TODO_USER: непредвиденный формат ответа, логировать при необходимости
+            text_value = str(first_item)
+    else:
+        text_value = str(result)
+
+    recognized_text: str = text_value.strip()
     return recognized_text
 
 
 # ============================
-# 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: суммаризация (английский)
+# 3. Суммаризация (английский DistilBART)
 # ============================
 
 summary_pipeline = pipeline(
@@ -121,8 +98,8 @@ def run_summarization(
         max(32, word_count + 20),
     )
 
+    # Для совсем короткого текста суммаризация мало смысла
     if word_count < 8:
-        # TODO_USER: для очень короткого текста суммаризация сомнительна, возвращаем исходный текст
         return cleaned_text
 
     summary_result_list = summary_pipeline(
@@ -137,7 +114,7 @@ def run_summarization(
 
 
 # ============================
-# 5. Трансформер #3: TTS (английский, MMS VITS)
+# 4. TTS (английский, MMS VITS)
 # ============================
 
 tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-eng")
@@ -150,7 +127,7 @@ def run_tts(summary_text: str) -> Optional[str]:
     Озвучка английского текста конспекта через VitsModel (facebook/mms-tts-eng).
 
     Если модель внутри упадёт (известный баг на некоторых странных инпутах),
-    мы просто вернём None и не будем ронять всё приложение.
+    просто возвращаем None и не роняем всё приложение.
     """
     cleaned_text: str = summary_text.strip()
     if not cleaned_text:
@@ -195,24 +172,21 @@ def run_tts(summary_text: str) -> Optional[str]:
 
 
 # ============================
-# 6. Полный пайплайн
+# 5. Полный пайплайн
 # ============================
 
 def full_flow(
     image_object: Image.Image,
     max_summary_tokens: int = 128,
-) -> Tuple[str, str, str, Optional[str]]:
+) -> Tuple[str, str, Optional[str]]:
     """
     Полный пайплайн:
-    1) OCR (easyocr): изображение -> исходный текст (английский)
-    2) Классификация текста трансформером (sentiment)
-    3) Суммаризация: текст -> конспект
-    4) TTS: конспект -> .wav файл (или None)
+    1) OCR: изображение -> исходный английский текст (olmOCR)
+    2) Суммаризация: текст -> конспект (DistilBART)
+    3) TTS: конспект -> .wav файл (или None, если TTS не смог)
     """
     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,
@@ -220,11 +194,11 @@ def full_flow(
 
     audio_file_path: Optional[str] = run_tts(summary_text=summary_text)
 
-    return recognized_text, classification_text, summary_text, audio_file_path
+    return recognized_text, summary_text, audio_file_path
 
 
 # ============================
-# 7. Gradio UI (на русском)
+# 6. Gradio UI (по-русски)
 # ============================
 
 gradio_interface = gradio_module.Interface(
@@ -244,29 +218,25 @@ gradio_interface = gradio_module.Interface(
     ],
     outputs=[
         gradio_module.Textbox(
-            label="Распознанный текст (OCR, easyocr)",
-            lines=8,
-        ),
-        gradio_module.Textbox(
-            label="Анализ текста (классификация, DistilBERT)",
-            lines=2,
+            label="Распознанный текст (olmOCR)",
+            lines=10,
         ),
         gradio_module.Textbox(
-            label="Конспект (английский текст, DistilBART)",
+            label="Конспект (английский текст)",
             lines=6,
         ),
         gradio_module.Audio(
-            label="Озвучка конспекта (английский TTS, VITS)",
+            label="Озвучка конспекта (английский TTS)",
             type="filepath",
         ),
     ],
-    title="Картинка → Текст → Анализ → Конспект → Озвучка",
+    title="Картинка → Текст → Конспект → Озвучка (olmOCR + английские модели)",
     description=(
-        "1) easyocr распознаёт печатный английский текст с картинки.\n"
-        "2) Трансформер-классификатор (DistilBERT) оценивает тон текста.\n"
-        "3) Трансформер-суммаризатор (DistilBART) делает краткий конспект.\n"
-        "4) Трансформер TTS (MMS VITS) озвучивает конспект.\n"
-        "В проекте используются три трансформера с Hugging Face, OCR сделан через easyocr."
+        "1) olmOCR-2-7B-1025-FP8 распознаёт текст с документа.\n"
+        "2) Английский трансформер суммаризации д��лает краткий пересказ.\n"
+        "3) VITS-модель MMS (facebook/mms-tts-eng) озвучивает конспект.\n\n"
+        "Если озвучка не сгенерировалась, значит конкретный текст не понравился TTS-модели "
+        "и она упала внутри — пайплайн просто пропустит аудио."
     ),
 )