conjurer/voice_recognition.py

import logging
import wave

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
# Replace with your API key
aai.settings.api_key = "aa9962f0088a449a9c4ab2361e96cc08"
discord.opus._load_default()
# VoiceState
# self_mute=False
# self_deaf=False
# self_stream=False
# suppress=False
# requested_to_speak_at=None
# channel=<VoiceChannel id=1285599336318632036
#           name='testing'
#           rtc_region=None
#           position=10
#           bitrate=64000
#           video_quality_mode=<VideoQualityMode.auto: 1>
#           user_limit=0
#           category_id=1084451744496496753
# URL of the file to transcribe
# You can also transcribe a local file by passing in a file path
# FILE_URL = './path/to/file.mWp3
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):

        self.username = str(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):
        self.transcript_file_present.close()
        self.name_file_past = self.file_name_present
        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)
        # send present to transcription

    def writeframes(self, pcmdata):
        logger.info(self.transcript_file_present)
        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()
        # send present to transcription


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):
        super().__init__()
        
        self.wavewriter = {}

    def on_user_connect(self, username):
        self.wavewriter[str(username)] = [WaveWriter(username), time.time_ns(), False]

    def on_user_disconnect(self,username):
        self.wavewriter[str(username)].cleanup()
        self.wavewriter.pop(str(username))

    def rotate_user(self,username):
        self.wavewriter[str(username)].rotate()

    def wants_opus(self) -> bool:
        return False

    def write(self, user, data) -> None:
        logger.info("User %s", user)
        if user:
            self.wavewriter[str(user.id)][0].writeframes(data.pcm)
            self.wavewriter[str(user.id)][1] = time.time_ns()
            self.wavewriter[str(user.id)][2] = 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):
        self.bot = bot
        self.wsink = SRBuffer()
        self.threads = []
        self.comm_queue = Queue()

    async def transcribe(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:
            print(f"Speaker {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)
            vc.listen(self.wsink)
            self.check_data.start()

    @tasks.loop(seconds=5)
    async def check_data(self):
        for item in self.wsink.wavewriter.values():
            logger.info(item[0])
            logger.info(item[1])
            logger.info(item[2])
            logger.info(time.time_ns())

    @commands.Cog.listener()
    async def on_voice_state_update(self, user, before, after):
        if before.channel is None:
            logger.info("User %s connected to channel %s", user, after.channel.name)
            self.wsink.on_user_connect(user.id)
        elif after.channel is  None:
            logger.info("User %s disconnected from channel %s", user, before.channel.name)
            operation = {
                "type": "user_cleanup",
                user : str(user.id)
            }
            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):
        await ctx.voice_client.disconnect()


async def setup(bot):
    await bot.add_cog(Transcriber(bot))


# 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.