itis / app.py
nktssk's picture
Update app.py
05776be verified
raw
history blame
7.87 kB
import os
import io
import torch
import gradio as gr
import wikipediaapi
import re
import inflect
import soundfile as sf
import unicodedata
import num2words
from PIL import Image
from datasets import load_dataset
from scipy.io.wavfile import write
from transformers import VitsModel, AutoTokenizer
from transformers import pipeline
from transformers import T5ForConditionalGeneration, T5Tokenizer
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from google.cloud import vision
from transformers import CLIPProcessor, CLIPModel
########################################
# (Опционально) Установите переменную окружения для Google Cloud:
# os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/path/to/your/service_account.json"
########################################
def clean_text(text):
# Очистка от некоторых спецсимволов, ссылок, диакритики
text = re.sub(r'МФА:?\s?\[.*?\]', '', text)
text = re.sub(r'\[.*?\]', '', text)
def remove_diacritics(char):
if unicodedata.category(char) == 'Mn':
return ''
return char
text = unicodedata.normalize('NFD', text)
text = ''.join(remove_diacritics(char) for char in text)
text = unicodedata.normalize('NFC', text)
text = re.sub(r'\s+', ' ', text)
text = re.sub(r'[^\w\s.,!?-]', '', text)
return text.strip()
from num2words import num2words
def number_to_russian_text(number):
try:
return num2words(number, lang='ru')
except NotImplementedError:
return "Ошибка: Не поддерживается преобразование для этого числа."
summarization_model = pipeline("summarization", model="facebook/bart-large-cnn")
wiki = wikipediaapi.Wikipedia("Nikita", "en")
embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
t2s_pipe = pipeline("text-to-speech", model="facebook/mms-tts-rus")
translator = pipeline("translation_en_to_ru", model="Helsinki-NLP/opus-mt-en-ru")
def text_to_speech(text, output_path="speech.wav"):
text = number_to_russian_text(text)
model = VitsModel.from_pretrained("facebook/mms-tts-rus")
tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
output = model(**inputs).waveform.squeeze().numpy()
sf.write(output_path, output, samplerate=model.config.sampling_rate)
return output_path
def fetch_wikipedia_summary(landmark):
page = wiki.page(landmark)
if page.exists():
return clean_text(page.summary)
else:
return "Found error!"
def recognize_landmark_google_cloud(image):
client = vision.ImageAnnotatorClient()
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
img_bytes = io.BytesIO()
image.save(img_bytes, format='PNG')
content = img_bytes.getvalue()
vision_image = vision.Image(content=content)
response = client.landmark_detection(image=vision_image)
landmarks = response.landmark_annotations
if landmarks:
return landmarks[0].description
else:
return "Unknown"
def tourist_helper_english(landmark):
wiki_text = fetch_wikipedia_summary(landmark)
if wiki_text == "Found error!":
return None
summarized_text = summarization_model(wiki_text, min_length=20, max_length=210)[0]["summary_text"]
audio_path = text_to_speech(summarized_text)
return audio_path
def process_image_google_cloud(image):
recognized = recognize_landmark_google_cloud(image)
print(f"[GoogleVision] Распознано: {recognized}")
audio_path = tourist_helper_english(recognized)
return audio_path
def process_text_google_cloud(landmark):
return tourist_helper_english(landmark)
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
text_inputs = clip_processor(
text=landmark_titles,
images=None,
return_tensors="pt",
padding=True
)
with torch.no_grad():
text_embeds = clip_model.get_text_features(**text_inputs)
text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
def recognize_landmark_clip(image):
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
image_inputs = clip_processor(images=image, return_tensors="pt")
with torch.no_grad():
image_embed = clip_model.get_image_features(**image_inputs)
image_embed = image_embed / image_embed.norm(p=2, dim=-1, keepdim=True)
similarity = (image_embed @ text_embeds.T).squeeze(0)
best_idx = similarity.argmax().item()
best_score = similarity[best_idx].item()
recognized_landmark = landmark_titles[best_idx]
return recognized_landmark, best_score
def tourist_helper_with_russian(landmark):
wiki_text = fetch_wikipedia_summary(landmark)
if wiki_text == "Found error!":
return None
print(wiki_text)
summarized_text = summarization_model(wiki_text, min_length=20, max_length=210)[0]["summary_text"]
print(summarized_text)
translated = translator(summarized_text, max_length=1000)[0]["translation_text"]
print(translated)
audio_path = text_to_speech(translated)
return audio_path
def process_image_clip(image):
recognized, score = recognize_landmark_clip(image)
print(f"[CLIP] Распознано: {recognized}, score={score:.2f}")
audio_path = tourist_helper_with_russian(recognized)
return audio_path
def process_text_clip(landmark):
return tourist_helper_with_russian(landmark)
with gr.Blocks() as demo:
gr.Markdown("## Две демки: Google Cloud Vision и CLIP (с переводом на русский)")
with gr.Tabs():
with gr.Tab("CLIP + Sum + Translate + T2S"):
gr.Markdown("### Распознавание (CLIP) и перевод на русский")
with gr.Row():
image_input_c = gr.Image(label="Загрузите фото", type="pil")
text_input_c = gr.Textbox(label="Или введите название")
audio_output_c = gr.Audio(label="Результатт")
with gr.Row():
btn_recognize_c = gr.Button("Распознать и перевести на русский")
btn_text_c = gr.Button("Поиск по тексту")
btn_recognize_c.click(
fn=process_image_clip,
inputs=image_input_c,
outputs=audio_output_c
)
btn_text_c.click(
fn=process_text_clip,
inputs=text_input_c,
outputs=audio_output_c
)
with gr.Tab("Google + Sum + T2S"):
gr.Markdown("### Распознавание достопримечательности (Google)")
with gr.Row():
image_input_g = gr.Image(label="Загрузите фото", type="pil")
text_input_g = gr.Textbox(label="Или введите название вручную")
audio_output_g = gr.Audio(label="Результат")
with gr.Row():
btn_recognize_g = gr.Button("Распознать и озвучить")
btn_text_g = gr.Button("Распознать по тексту и озвучить")
btn_recognize_g.click(
fn=process_image_google_cloud,
inputs=image_input_g,
outputs=audio_output_g
)
btn_text_g.click(
fn=process_text_google_cloud,
inputs=text_input_g,
outputs=audio_output_g
)
demo.launch(debug=True)