models/censor.py

from utils.bad_words import words
import re
from models.whisper import model
import random
import string
import subprocess

class Wash:
    def __init__(self, custom_db = None):
        self.db = words
        if(custom_db):
          self.db = custom_db

    @staticmethod
    def clean(phrase):
        return re.sub(r'\s+', ' ', phrase.lower())

    @staticmethod
    def tokenize(phrase):
        with_punctuation = re.sub(r'[^a-zA-Z\s]', '', phrase).split()
        without_punctuation = re.sub(r'[^a-zA-Z\s]', '', phrase).split()
        return with_punctuation + without_punctuation

    def words(self):
        return self.db

    def check_word(self, word):
      word = word.replace(" ", "")
      bad_words = self.words()
      tokens = self.tokenize(word)

      for token in tokens:
            if token in bad_words:
                return True

      return False
  
    def generate_random_name(self, length):
        letters = string.ascii_lowercase
        return ''.join(random.choice(letters) for _ in range(length))

    def insert_beep(self, itensList, original_file, offset=0):
        if not itensList:
            return
        output = self.generate_random_name(8) + ".wav"

        filter_str = "[0]volume=0:enable=\'"
        for i, item in enumerate(itensList):
            start = item["start"]
            end = item["end"]
            start_plus_offset =  round(start + offset, 2)
            end_minus_offset =  round(end - offset, 2)
            filter_str += f'between(t,{start_plus_offset},{end_minus_offset})'

            if i != len(itensList) - 1:
                filter_str += '+'
            else:
                filter_str += "\'"
        for item in itensList:
            start = round(item["start"], 2)
            end = round(item["end"], 2)
            start_plus_offset = round(start + offset, 2)
            end_minus_offset = round(end - offset, 2)

            filter_str += f'[main{start}];'
            filter_str += f'sine=d={round(end_minus_offset - start_plus_offset, 2)}:f=1000,adelay={start_plus_offset * 1000}s,pan=stereo|FL=c0|FR=c0[beep{start}];'
            filter_str += f'[main{start}][beep{start}]amix=inputs=2:duration=longest'

        filter_str = filter_str.rstrip(',')

        ffmpeg_cmd = [
            "ffmpeg",
            "-i",
            original_file,
            "-filter_complex",
            filter_str,
            output
        ]

        result = subprocess.run(ffmpeg_cmd)

        return output
    
    def process_audio(self, audio_path, offset):
        bad_time = []
        segments, info = model.transcribe(audio_path, word_timestamps=True, beam_size=1, best_of=1)
        for segment in segments:
          for word in segment.words:
              is_profane = self.check_word(word.word)
              if(is_profane):
                bad_time.append({
                    "start": word.start,
                    "end": word.end,
                })

        result = self.insert_beep(bad_time, audio_path, offset)
        if(result):
          return result
        return audio_path

    def check(self, phrase):
        cleaned_phrase = self.clean(phrase)
        cleaned_phrase = phrase.replace(" ", "")
        tokens = self.tokenize(cleaned_phrase)
        bad_words = self.words()

        for token in tokens:
            if token in bad_words:
                return True

        return False