edit

app.py

@@ -1,181 +1,83 @@
 from typing import Tuple, Optional
-
 import tempfile
-
-import numpy as np
-import soundfile as soundfile_module
+import soundfile as sf
 import torch
-import gradio as gradio_module
+import gradio as gr
+import numpy as np
 from PIL import Image
 from transformers import (
     TrOCRProcessor,
     VisionEncoderDecoderModel,
     pipeline,
-    VitsModel,
     AutoTokenizer,
+    VitsModel,
 )
 
-ocr_processor: TrOCRProcessor = TrOCRProcessor.from_pretrained(
-    "microsoft/trocr-small-printed"
-)
-ocr_model: VisionEncoderDecoderModel = VisionEncoderDecoderModel.from_pretrained(
-    "microsoft/trocr-small-printed"
-)
+
+ocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
+ocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
+ocr_model.to("cpu")
 
 summary_pipeline = pipeline(
-    task="summarization",
-    model="sshleifer/distilbart-cnn-12-6",
+    "summarization",
+    model="IlyaGusev/mbart_ru_sum_gazeta",
+    tokenizer="IlyaGusev/mbart_ru_sum_gazeta",
 )
 
-tts_model: VitsModel = VitsModel.from_pretrained("facebook/mms-tts-rus")
-tts_tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
-
-device_string: str = "cpu"
-ocr_model.to(device_string)
-tts_model.to(device_string)
+tts_model = VitsModel.from_pretrained("facebook/mms-tts-rus")
+tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
+tts_model.to("cpu")
 
-
-def run_ocr(image_object: Image.Image) -> str:
-    """
-    Распознавание текста с изображения.
-    Предполагаем, что на картинке простой напечатанный текст.
-    """
-    if image_object is None:
+def run_ocr(image: Image.Image) -> str:
+    if image is None:
         return ""
-
-    processor_output = ocr_processor(
-        images=image_object,
-        return_tensors="pt",
-    )
-    pixel_values_tensor = processor_output.pixel_values.to(device_string)
-
-    generated_id_tensor = ocr_model.generate(pixel_values_tensor)
-    decoded_text_list = ocr_processor.batch_decode(
-        generated_id_tensor,
-        skip_special_tokens=True,
-    )
-
-    recognized_text: str = decoded_text_list[0]
-    return recognized_text.strip()
-
-
-def run_summarization(
-    input_text: str,
-    max_summary_tokens: int = 128,
-) -> str:
-    """
-    Суммаризация текста до короткого конспекта.
-    Без сложного разбиения на чанки -> длинные тексты лучше не кормить.
-    """
-    cleaned_text: str = input_text.strip()
-    if not cleaned_text:
+    pixel_values = ocr_processor(images=image, return_tensors="pt").pixel_values
+    generated_ids = ocr_model.generate(pixel_values)
+    text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return text.strip()
+
+def run_summary(text: str) -> str:
+    text = text.strip()
+    if not text:
         return ""
+    result = summary_pipeline(text, max_length=128, min_length=30, do_sample=False)
+    return result[0]["summary_text"].strip()
 
-    word_count: int = len(cleaned_text.split())
-    dynamic_max_length: int = min(
-        max_summary_tokens,
-        max(32, word_count + 20),
-    )
-
-    summary_result_list = summary_pipeline(
-        cleaned_text,
-        max_length=dynamic_max_length,
-        min_length=max(10, dynamic_max_length // 3),
-        do_sample=False,
-    )
-
-    summary_text: str = summary_result_list[0]["summary_text"].strip()
-    return summary_text
-
-
-def run_tts(summary_text: str) -> Optional[str]:
-    """
-    Озвучка текста конспекта через VitsModel (facebook/mms-tts-rus).
-    Возвращаем путь до временного .wav файла, который Gradio отдаст в плеер.
-    """
-    cleaned_text: str = summary_text.strip()
-    if not cleaned_text:
+def run_tts(text: str) -> Optional[str]:
+    text = text.strip()
+    if not text:
         return None
-
-    tokenized_inputs = tts_tokenizer(
-        cleaned_text,
-        return_tensors="pt",
-    ).to(device_string)
-
+    inputs = tts_tokenizer(text, return_tensors="pt").to("cpu")
     with torch.no_grad():
-        model_output = tts_model(**tokenized_inputs)
-        waveform_tensor = model_output.waveform
-
-    waveform_array = waveform_tensor.squeeze().cpu().numpy().astype("float32")
-
-    with tempfile.NamedTemporaryFile(
-        suffix=".wav",
-        delete=False,
-    ) as temporary_file:
-        soundfile_module.write(
-            temporary_file.name,
-            waveform_array,
-            tts_model.config.sampling_rate,
-        )
-        file_path: str = temporary_file.name
+        waveform = tts_model(**inputs).waveform
+    audio = waveform.squeeze().cpu().numpy().astype("float32")
+    audio = np.clip(audio, -1.0, 1.0)
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        sf.write(f.name, audio, tts_model.config.sampling_rate)
+        return f.name
 
-    return file_path
+def full_flow(image: Image.Image) -> Tuple[str, str, Optional[str]]:
+    text = run_ocr(image)
+    summary = run_summary(text)
+    audio_path = run_tts(summary)
+    return text, summary, audio_path
 
 
-def full_flow(
-    image_object: Image.Image,
-    max_summary_tokens: int = 128,
-) -> Tuple[str, str, Optional[str]]:
-
-    recognized_text: str = run_ocr(image_object=image_object)
-
-    summary_text: str = run_summarization(
-        input_text=recognized_text,
-        max_summary_tokens=max_summary_tokens,
-    )
-
-    audio_file_path: Optional[str] = run_tts(summary_text=summary_text)
-
-    return recognized_text, summary_text, audio_file_path
-
-
-gradio_interface = gradio_module.Interface(
+demo = gr.Interface(
     fn=full_flow,
-    inputs=[
-        gradio_module.Image(
-            type="pil",
-            label="Изображение с напечатанным текстом (лучше русским/латиницей)",
-        ),
-        gradio_module.Slider(
-            minimum=32,
-            maximum=256,
-            value=128,
-            step=16,
-            label="Максимальная длина конспекта (токены, примерно)",
-        ),
-    ],
+    inputs=gr.Image(type="pil", label="Изображение с текстом (русский или английский)"),
     outputs=[
-        gradio_module.Textbox(
-            label="Распознанный текст (OCR)",
-            lines=6,
-        ),
-        gradio_module.Textbox(
-            label="Конспект (суммаризация)",
-            lines=6,
-        ),
-        gradio_module.Audio(
-            label="Озвучка конспекта (VITS, ru)",
-            type="filepath",
-        ),
+        gr.Textbox(label="Распознанный текст", lines=6),
+        gr.Textbox(label="Краткий пересказ", lines=6),
+        gr.Audio(label="Озвучка конспекта", type="filepath"),
     ],
-    title="Картинка → Конспект → Озвучка (Transformers)",
+    title="Картинка → Текст → Конспект → Озвучка (русская версия)",
     description=(
-        "1) Трансформер OCR распознаёт текст с изображения. "
-        "2) Трансформер суммаризации сокращает текст до конспекта. "
-        "3) VITS-модель (facebook/mms-tts-rus) озвучивает конспект по-русски."
+        "1️⃣ OCR (TrOCR-base) распознаёт текст с картинки.\n"
+        "2️⃣ Суммаризация (IlyaGusev/mbart_ru_sum_gazeta) делает конспект.\n"
+        "3️⃣ TTS (facebook/mms-tts-rus) озвучивает результат."
     ),
 )
 
-
 if __name__ == "__main__":
-    gradio_interface.launch()
+    demo.launch()