mirror of
https://github.com/migatu/conjurer.git
synced 2026-07-14 21:38:38 +00:00
232 lines
9.0 KiB
Python
232 lines
9.0 KiB
Python
import logging
|
|
import wave
|
|
import asyncio
|
|
import assemblyai as aai
|
|
import discord
|
|
from discord.ext import commands, voice_recv, tasks
|
|
from discord.opus import Decoder as OpusDecoder
|
|
from queue import Queue, Empty
|
|
import time
|
|
import threading
|
|
# Replace with your API key
|
|
aai.settings.api_key = "aa9962f0088a449a9c4ab2361e96cc08"
|
|
discord.opus._load_default()
|
|
CHANNELS = OpusDecoder.CHANNELS
|
|
SAMPLE_WIDTH = OpusDecoder.SAMPLE_SIZE // OpusDecoder.CHANNELS
|
|
SAMPLING_RATE = OpusDecoder.SAMPLING_RATE
|
|
|
|
#rotate file after there is 0.5s between last received pcm for user.
|
|
#delete messsages after user disconnect
|
|
|
|
logger = logging.getLogger("discord")
|
|
location = "/home/pi/Conjurer/transcripts/"
|
|
|
|
|
|
class WaveWriter:
|
|
def __init__(self, user_id, queue):
|
|
self.queue = queue
|
|
self.user = user_id
|
|
self.username = str(user_id.id)
|
|
logger.info("Creating Wavewriter %s", self.username)
|
|
logger.info(user_id)
|
|
self.present_file_id = 0
|
|
self.file_name_past = None
|
|
self.file_name_present = location + self.username + "_" + str(self.present_file_id) + ".mp3"
|
|
self.transcript_file_present : wave.Wave_write = wave.open(
|
|
self.file_name_present, "wb"
|
|
)
|
|
self.transcript_file_present.setnchannels(CHANNELS)
|
|
self.transcript_file_present.setsampwidth(SAMPLE_WIDTH)
|
|
self.transcript_file_present.setframerate(SAMPLING_RATE)
|
|
self.file_name_future = location + self.username + "_" + str(self.present_file_id + 1) + ".mp3"
|
|
self.transcript_file_future : wave.Wave_write = wave.open(
|
|
self.file_name_future, "wb"
|
|
)
|
|
self.transcript_file_future.setnchannels(CHANNELS)
|
|
self.transcript_file_future.setsampwidth(SAMPLE_WIDTH)
|
|
self.transcript_file_future.setframerate(SAMPLING_RATE)
|
|
|
|
def rotate(self):
|
|
|
|
logger.info("Before rotation")
|
|
logger.info(self.file_name_present)
|
|
logger.info(self.file_name_past)
|
|
|
|
self.transcript_file_present.close()
|
|
self.file_name_past = self.file_name_present
|
|
if self.present_file_id > 10:
|
|
self.present_file_id = 0
|
|
else:
|
|
self.present_file_id += 1
|
|
self.file_name_present = self.file_name_future
|
|
self.transcript_file_present = self.transcript_file_future
|
|
self.file_name_future = location + self.username + "_" + str(self.present_file_id + 1) + ".mp3"
|
|
self.transcript_file_future : wave.Wave_write = wave.open(
|
|
self.file_name_future, "wb"
|
|
)
|
|
self.transcript_file_future.setnchannels(CHANNELS)
|
|
self.transcript_file_future.setsampwidth(SAMPLE_WIDTH)
|
|
self.transcript_file_future.setframerate(SAMPLING_RATE)
|
|
operation = {
|
|
"type": "user_cleanup",
|
|
"user" : self.user,
|
|
"filename": self.file_name_present
|
|
}
|
|
logger.info("After rotation")
|
|
logger.info(self.file_name_present)
|
|
logger.info(self.file_name_past)
|
|
|
|
self.queue.put(operation)
|
|
|
|
def writeframes(self, pcmdata):
|
|
self.transcript_file_present.writeframes(pcmdata)
|
|
|
|
def cleanup(self):
|
|
logger.info("Cleanup for user %s", self.username)
|
|
self.transcript_file_present.close()
|
|
self.transcript_file_future.close()
|
|
|
|
|
|
class SRBuffer(voice_recv.AudioSink):
|
|
"""Endpoint AudioSink that generates a wav file.
|
|
Best used in conjunction with a silence generating sink. (TBD)
|
|
"""
|
|
# on member join dodajemytypa do listy
|
|
# on member disconnect - dropujemy go
|
|
def __init__(self, queue):
|
|
super().__init__()
|
|
self.queue = queue
|
|
self.wavewriter = {}
|
|
|
|
def on_user_connect(self, username):
|
|
self.wavewriter[str(username.id)] = [WaveWriter(username, self.queue), time.time_ns(), time.time_ns(), False, False]
|
|
|
|
def on_user_disconnect(self,username):
|
|
self.wavewriter[str(username.id)].cleanup()
|
|
self.wavewriter.pop(str(username.id))
|
|
|
|
def wants_opus(self) -> bool:
|
|
return False
|
|
|
|
def write(self, user, data) -> None:
|
|
if user:
|
|
self.wavewriter[str(user.id)][0].writeframes(data.pcm)
|
|
self.wavewriter[str(user.id)][1] = time.time_ns() #time from last write
|
|
self.wavewriter[str(user.id)][2] = time.time_ns() #time since last rotation
|
|
self.wavewriter[str(user.id)][3] = True
|
|
self.wavewriter[str(user.id)][4] = True
|
|
|
|
def cleanup(self) -> None:
|
|
try:
|
|
logger.info("Cleanup for SRBuffer")
|
|
for item in self.wavewriter.values():
|
|
item[0].cleanup()
|
|
except Exception:
|
|
logger.warning("WaveSink got error closing file on cleanup", exc_info=True)
|
|
|
|
|
|
class Transcriber(commands.Cog):
|
|
def __init__(self, bot, worker, queue):
|
|
self.bot = bot
|
|
self.threads = []
|
|
self.comm_queue = queue
|
|
self.worker = worker
|
|
self.wsink = SRBuffer(self.comm_queue)
|
|
async def transcribe(self,file_url):
|
|
config = aai.TranscriptionConfig(speaker_labels=True, language_code="pl")
|
|
|
|
transcriber = aai.Transcriber()
|
|
transcript = transcriber.transcribe(file_url, config=config)
|
|
for utterance in transcript.utterances:
|
|
logger.info("Speaker %s : %s", utterance.speaker, utterance.text)
|
|
|
|
@commands.hybrid_command(name="transcribe")
|
|
async def test(self, ctx):
|
|
logger.info("Attempt transcribe")
|
|
vc = None
|
|
|
|
if self.bot.voice_clients:
|
|
if isinstance(self.bot.voice_clients[0], voice_recv.VoiceRecvClient):
|
|
logger.info("Already transcribing")
|
|
else:
|
|
logger.info("Already connected with other client")
|
|
logger.info(self.bot.voice_clients)
|
|
else:
|
|
logger.info("Connected")
|
|
vc = await ctx.author.voice.channel.connect(cls=voice_recv.VoiceRecvClient)
|
|
logger.info(self.bot.voice_clients)
|
|
self.output_handler.work = True
|
|
self.output_handler.scan_loop.start()
|
|
self.worker.start()
|
|
vc.listen(self.wsink)
|
|
|
|
@tasks.loop(seconds=0.5)
|
|
async def check_data(self):
|
|
for item in self.wsink.wavewriter.values():
|
|
logger.info("Item %s", item)
|
|
if not item[3]:
|
|
timediff_rotation = time.time_ns() - item[2]
|
|
timediff_write = time.time_ns() -item[1]
|
|
if timediff_rotation > 3000149433 and timediff_write > 500014943 and item[4]:
|
|
logger.info("File rotation time since last write %s", timediff_write/1e9)
|
|
logger.info("File rotation time since last rotation %s", timediff_rotation/1e9)
|
|
item[4] = False
|
|
item[0].rotate()
|
|
#temporary - will be put in separate thread
|
|
await self.transcribe(item[0].file_name_past)
|
|
item[4] = False
|
|
logger.info("Diff : %s", time.time_ns() - item[1])
|
|
|
|
@tasks.loop(seconds=1)
|
|
async def scan_queue(self):
|
|
self.comm_queue.get()
|
|
|
|
@commands.Cog.listener()
|
|
async def on_voice_state_update(self, user, before, after):
|
|
if user == self.bot.user:
|
|
logger.info("Ignoring self")
|
|
return
|
|
|
|
if before.channel is None:
|
|
logger.info("User %s connected to channel %s", user, after.channel.name)
|
|
self.wsink.on_user_connect(user)
|
|
elif after.channel is None:
|
|
logger.info("User %s disconnected from channel %s", user, before.channel.name)
|
|
operation = {
|
|
"type": "user_cleanup",
|
|
"user" : user,
|
|
"filename": None
|
|
}
|
|
self.comm_queue.put(operation)
|
|
else:
|
|
logger.info("User VC status changed %s", user.id)
|
|
logger.info("Before %s", before)
|
|
logger.info("After %s", after)
|
|
|
|
|
|
@commands.command(name="stop_transcribe")
|
|
async def stop(self, ctx):
|
|
|
|
self.output_handler.work = False
|
|
self.check_data.stop()
|
|
self.worker.stop()
|
|
await ctx.voice_client.disconnect()
|
|
|
|
def transcribe_output_queue(queue):
|
|
while True:
|
|
queue.get()
|
|
logger.info("PING")
|
|
|
|
async def setup(bot):
|
|
logger = logging.getLogger("discord")
|
|
logger.info("Loading voice transcribed module phase one")
|
|
worker = threading.Thread(target=transcribe_output_queue)
|
|
logger.info("Loading voice transcribed phase two")
|
|
queue = Queue()
|
|
logger.info("Loading voice transcribed module phase three")
|
|
await bot.add_cog(Transcriber(bot, worker, queue))
|
|
logger.info("Loading voice transcribed module done")
|
|
|
|
# 1. zapisuj kwestie człowieka do pliku w którym będzie można stwierdzić kto co powiedział (callback z basicaudio + write z wavesinka). Zamykaj plik i wysyłaj do transkrypcji w momencie ciszy dłuższej niż 0.5s. Jeśli człowiek nadaje cały czas dawaj sygnał że nie można go transkrybować - i wyłaczaj zapis.
|
|
# 2. Transkrypcja to abstrakt - w zależności od tego która metoda jest włączona wysyła do odpowiedniego silnika, silniki offline odpalam na activcomie. Zaczynamy od assemblyai bo najlepiej wspiera polski. Potem zobaczymy co dalej.
|