ASureevaA commited on
Commit
8dac34d
·
1 Parent(s): a88eb1e
Files changed (2) hide show
  1. app.py +14 -26
  2. requirements.txt +0 -2
app.py CHANGED
@@ -1,10 +1,8 @@
1
  from typing import Tuple, Optional, Any
2
 
3
- import torch
4
  import numpy as numpy
5
- import gradio as gr
6
  from PIL import Image
7
- from datasets import load_dataset
8
  from transformers import (
9
  TrOCRProcessor,
10
  VisionEncoderDecoderModel,
@@ -23,22 +21,15 @@ summary_pipeline = pipeline(
23
  model="sshleifer/distilbart-cnn-12-6",
24
  )
25
 
26
- tts_pipeline = pipeline(
27
  task="text-to-speech",
28
- model="microsoft/speecht5_tts",
29
  )
30
 
31
- speaker_dataset = load_dataset(
32
- path="Matthijs/cmu-arctic-xvectors",
33
- split="validation",
34
- )
35
- speaker_embedding_tensor: torch.Tensor = torch.tensor(
36
- speaker_dataset[7306]["xvector"]
37
- ).unsqueeze(0)
38
 
39
  def run_ocr(image_object: Image.Image) -> str:
40
  """
41
- Распознавание текста с изображения с помощью трансформера OCR.
42
  Предполагаем, что на картинке простой напечатанный текст.
43
  """
44
  if image_object is None:
@@ -65,8 +56,8 @@ def run_summarization(
65
  max_summary_tokens: int = 128,
66
  ) -> str:
67
  """
68
- Суммаризация текста.
69
- Здесь без разбиения на чанки, поэтому для очень длинных текстов могут быть проблемы.
70
  """
71
  cleaned_text: str = input_text.strip()
72
  if not cleaned_text:
@@ -86,16 +77,13 @@ def run_summarization(
86
  def run_tts(summary_text: str) -> Optional[Tuple[int, Any]]:
87
  """
88
  Озвучка текста конспекта.
89
- Возвращает кортеж (частота_дискретизации, аудиоданные) или None, если текста нет.
90
  """
91
  cleaned_text: str = summary_text.strip()
92
  if not cleaned_text:
93
  return None
94
 
95
- tts_output = tts_pipeline(
96
- cleaned_text,
97
- forward_params={"speaker_embeddings": speaker_embedding_tensor},
98
- )
99
 
100
  sampling_rate_int: int = int(tts_output["sampling_rate"])
101
  audio_array = tts_output["audio"]
@@ -122,14 +110,14 @@ def full_flow(
122
 
123
  return recognized_text, summary_text, audio_tuple
124
 
125
- gradio_interface = gr.Interface(
126
  fn=full_flow,
127
  inputs=[
128
- gr.Image(
129
  type="pil",
130
  label="Изображение с напечатанным текстом (английский)",
131
  ),
132
- gr.Slider(
133
  minimum=32,
134
  maximum=256,
135
  value=128,
@@ -138,15 +126,15 @@ gradio_interface = gr.Interface(
138
  ),
139
  ],
140
  outputs=[
141
- gr.Textbox(
142
  label="Распознанный текст (OCR)",
143
  lines=6,
144
  ),
145
- gr.Textbox(
146
  label="Конспект (суммаризация)",
147
  lines=6,
148
  ),
149
- gr.Audio(
150
  label="Озвучка конспекта (TTS)",
151
  type="numpy",
152
  ),
 
1
  from typing import Tuple, Optional, Any
2
 
 
3
  import numpy as numpy
4
+ import gradio as gradio_module
5
  from PIL import Image
 
6
  from transformers import (
7
  TrOCRProcessor,
8
  VisionEncoderDecoderModel,
 
21
  model="sshleifer/distilbart-cnn-12-6",
22
  )
23
 
24
+ text_to_speech_pipeline = pipeline(
25
  task="text-to-speech",
26
+ model="facebook/mms-tts-eng",
27
  )
28
 
 
 
 
 
 
 
 
29
 
30
  def run_ocr(image_object: Image.Image) -> str:
31
  """
32
+ Распознавание текста с изображения.
33
  Предполагаем, что на картинке простой напечатанный текст.
34
  """
35
  if image_object is None:
 
56
  max_summary_tokens: int = 128,
57
  ) -> str:
58
  """
59
+ Суммаризация текста до короткого конспекта.
60
+ Без разбиения на чанки, поэтому огромные тексты лучше не подавать.
61
  """
62
  cleaned_text: str = input_text.strip()
63
  if not cleaned_text:
 
77
  def run_tts(summary_text: str) -> Optional[Tuple[int, Any]]:
78
  """
79
  Озвучка текста конспекта.
80
+ Используем модель, которой не нужны внешние speaker embeddings.
81
  """
82
  cleaned_text: str = summary_text.strip()
83
  if not cleaned_text:
84
  return None
85
 
86
+ tts_output = text_to_speech_pipeline(cleaned_text)
 
 
 
87
 
88
  sampling_rate_int: int = int(tts_output["sampling_rate"])
89
  audio_array = tts_output["audio"]
 
110
 
111
  return recognized_text, summary_text, audio_tuple
112
 
113
+ gradio_interface = gradio_module.Interface(
114
  fn=full_flow,
115
  inputs=[
116
+ gradio_module.Image(
117
  type="pil",
118
  label="Изображение с напечатанным текстом (английский)",
119
  ),
120
+ gradio_module.Slider(
121
  minimum=32,
122
  maximum=256,
123
  value=128,
 
126
  ),
127
  ],
128
  outputs=[
129
+ gradio_module.Textbox(
130
  label="Распознанный текст (OCR)",
131
  lines=6,
132
  ),
133
+ gradio_module.Textbox(
134
  label="Конспект (суммаризация)",
135
  lines=6,
136
  ),
137
+ gradio_module.Audio(
138
  label="Озвучка конспекта (TTS)",
139
  type="numpy",
140
  ),
requirements.txt CHANGED
@@ -1,8 +1,6 @@
1
  transformers
2
  torch
3
- datasets
4
  sentencepiece
5
- soundfile
6
  gradio
7
  Pillow
8
  numpy
 
1
  transformers
2
  torch
 
3
  sentencepiece
 
4
  gradio
5
  Pillow
6
  numpy