From ed0db782420b14e933ee39ff639c59de8609920d Mon Sep 17 00:00:00 2001 From: Hykilpikonna Date: Wed, 22 Dec 2021 13:13:21 -0500 Subject: [PATCH] [+] Telegram bot --- .gitignore | 3 ++ requirements.txt | 2 + src/bot.py | 113 +++++++++++++++++++++++++++++++++++++++++++++++ src/ina_main.py | 48 +++++++++++--------- 4 files changed, 144 insertions(+), 22 deletions(-) create mode 100644 src/bot.py diff --git a/.gitignore b/.gitignore index aeae209..e002528 100644 --- a/.gitignore +++ b/.gitignore @@ -117,3 +117,6 @@ dmypy.json # Custom .idea ._* +voice-bot-token.txt +audio_tmp +src/voice-bot-token.txt diff --git a/requirements.txt b/requirements.txt index a9801e3..86b6585 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,5 @@ tensorflow==2.2.0 plaidml-keras==0.7.0 inaSpeechSegmenter==0.6.8 + +python-telegram-bot diff --git a/src/bot.py b/src/bot.py new file mode 100644 index 0000000..f484890 --- /dev/null +++ b/src/bot.py @@ -0,0 +1,113 @@ +import os +from datetime import datetime +from pathlib import Path +from telegram import Update, Message +from telegram.ext import Updater, CallbackContext, Dispatcher, CommandHandler, MessageHandler, \ + Filters + +os.environ['KERAS_BACKEND'] = "plaidml.keras.backend" + +import keras + +from inaSpeechSegmenter import * +from ina_main import * +import warnings + +warnings.filterwarnings("ignore") + + +def r(u: Update, msg: str, md=True): + updater.bot.sendMessage(chat_id=u.effective_chat.id, text=msg, + parse_mode='Markdown' if md else None) + + +def cmd_start(u: Update, c: CallbackContext): + r(u, '欢迎! 点下面的录音按钮就可以开始啦w') + + +def process_audio(message: Message): + # Only when replying to voice or audio + audio = message.audio or message.voice + if not audio: + return + + # Download audio file + date = datetime.now().strftime('%Y-%m-%d %H-%M') + downloader = bot.getFile(audio.file_id) + file = Path(tmpdir).joinpath(f'{date} {message.from_user.name[1:]}.mp3') + print(downloader, '->', file) + downloader.download(file) + + # Segment file + seg = Segmenter() + result = process(seg, [str(file.absolute())]).results[0] + + # Null case + print(result.frames) + if len(result.frames) == 0: + bot.send_message(message.chat_id, '分析失败, 大概是音量太小或者时长太短吧, 再试试w') + return + + # Draw results + with draw_result(str(file), result) as buf: + f, m, o = get_result_percentages(result) + msg = f"分析结果: {f*100:.0f}% 🙋‍♀️ | {m*100:.0f}% 🙋‍♂️ | {o*100:.0f}% 🚫" + bot.send_photo(message.chat_id, photo=buf, caption=msg, + reply_to_message_id=message.message_id) + + +def cmd_analyze(u: Update, c: CallbackContext): + reply = u.effective_message.reply_to_message + + # Parse command + text = u.effective_message.text + if not text: + return + cmd = text.lower().split()[0].strip() + + if cmd[0] not in '!/': + return + cmd = cmd[1:] + + if cmd not in ['analyze', 'analyze-raw']: + return + + if cmd == 'analyze-raw': + raw = True + + if u.effective_user.id == reply.from_user.id: + process_audio(reply) + else: + r(u, '只有自己能分析自己的音频哦 👀') + + +def on_audio(u: Update, c: CallbackContext): + process_audio(u.effective_message) + + +if __name__ == '__main__': + tmpdir = Path('audio_tmp') + tmpdir.mkdir(exist_ok=True, parents=True) + + # Find telegram token + path = Path(os.path.abspath(__file__)).parent + db_path = path.joinpath('voice-bot-db.json') + if 'tg_token' in os.environ: + tg_token = os.environ['tg_token'] + else: + with open(path.joinpath('voice-bot-token.txt'), 'r', encoding='utf-8') as f: + tg_token = f.read().strip() + + # Telegram login + updater = Updater(token=tg_token, use_context=True) + dispatcher: Dispatcher = updater.dispatcher + bot = updater.bot + + dispatcher.add_handler(CommandHandler('start', cmd_start, filters=Filters.chat_type.private)) + dispatcher.add_handler(CommandHandler('analyze', cmd_analyze, filters=Filters.reply)) + dispatcher.add_handler(MessageHandler(Filters.reply, cmd_analyze)) + dispatcher.add_handler(MessageHandler(Filters.voice & Filters.chat_type.private, on_audio)) + dispatcher.add_handler(MessageHandler(Filters.audio & Filters.chat_type.private, on_audio)) + + print('Starting bot...') + updater.start_polling() diff --git a/src/ina_main.py b/src/ina_main.py index 92b2e9d..da3c1bf 100644 --- a/src/ina_main.py +++ b/src/ina_main.py @@ -2,6 +2,7 @@ from __future__ import annotations import io import os +import shutil import sys import tempfile import time @@ -62,7 +63,7 @@ def process(self: Segmenter, inp: list[str], tmpdir=None, verbose=False, skip_if break mspec, loge, diff_len = feats lseg = self.segment_feats(mspec, loge, diff_len, 0) - results.append(Result(lseg, inp[len(lmsg) - 1])) + results.append(Result([ResultFrame(*s) for s in lseg], inp[len(lmsg) - 1])) t_batch_dur = time.time() - t_batch_start nb_processed = len([e for e in lmsg if e[1] == 0]) @@ -93,6 +94,8 @@ def to_wav(file: str, callback: Callable, start_sec: float = 0, stop_sec: float output, error = p.communicate() assert p.returncode == 0, error + shutil.copy2(tmp_wav, './test_leohearts.wav') + return callback(tmp_wav) @@ -130,7 +133,7 @@ def draw_result(file: str, result: Result): ax.set_xlim([result.frames[0].start, result.frames[-1].end]) # Draw segmentation areas - colors = {'female': '#F5A9B8', 'male': '#5BCEFA', 'default': 'white'} + colors = {'female': '#F5A9B8', 'male': '#5BCEFA', 'default': 'gray'} for r in result.frames: color = colors[r.gender] if r.gender in colors else colors['default'] ax.axvspan(r.start, r.end - 0.01, alpha=.5, color=color) @@ -140,6 +143,8 @@ def draw_result(file: str, result: Result): plt.axis('off') plt.savefig(buf, bbox_inches='tight', pad_inches=0, transparent=False) buf.seek(0) + plt.clf() + plt.close() return buf return to_wav(file, wav_callback) @@ -165,27 +170,26 @@ def get_result_percentages(result: Result) -> tuple[float, float, float]: durations[d] /= total_dur # Return results - f = durations['female'] - m = durations['male'] + f = durations.get('female', 0) + m = durations.get('male', 0) return f, m, 1 - f - m +# def test(): +# results: BatchResults = BatchResults( +# [Result([ResultFrame('female', 0.0, 10.48), ResultFrame('male', 10.48, 12.780000000000001)], +# '../test.csv')], +# 1.7032792568206787, 1.7032792568206787, 1, +# [('../test.csv', 0)]) +# +# with draw_result('../test.mp3', results.results[0]) as buf: +# show_image_buffer(buf) +# print(get_result_percentages(results.results[0])) +# +# # seg = Segmenter() +# # print(process(seg, ['../test.mp3'])) + if __name__ == '__main__': - args = sys.argv[1:] - if len(args) < 0: - exit(0) - - results: BatchResults = BatchResults( - [Result([ResultFrame('female', 0.0, 10.48), ResultFrame('male', 10.48, 12.780000000000001)], - '../test.csv')], - 1.7032792568206787, 1.7032792568206787, 1, - [('../test.csv', 0)]) - - with draw_result('../test.mp3', results.results[0]) as buf: - show_image_buffer(buf) - print(get_result_percentages(results.results[0])) - - # inp = args[0] - - # seg = Segmenter() - # print(process(seg, ['../test.mp3'])) + to_wav('../audio_tmp/2021-12-22 05-32 leph1art5.mp3', print) + # test() + pass