app.py

import gradio as gr
import json
from difflib import Differ
import ffmpeg
import os
from pathlib import Path
import time
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import MarianMTModel, MarianTokenizer
import pandas as pd
import re
import time
import os
from fuzzywuzzy import fuzz
from fastT5 import export_and_get_onnx_model
import torch
from transformers import pipeline

MODEL  = "Finnish-NLP/wav2vec2-base-fi-voxpopuli-v2-finetuned"
marian_nmt_model = "Helsinki-NLP/opus-mt-tc-big-fi-en"
tokenizer_marian = MarianTokenizer.from_pretrained(marian_nmt_model)
model = MarianMTModel.from_pretrained(marian_nmt_model)

cuda = torch.device(
    'cuda:0') if torch.cuda.is_available() else torch.device('cpu')
sr_pipeline_device = 0 if torch.cuda.is_available() else -1

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
speech_recognizer = pipeline(
    task="automatic-speech-recognition",
    model=f'{MODEL}',
    tokenizer=f'{MODEL}',
    framework="pt",
    device=sr_pipeline_device,
)    

model_checkpoint = 'Finnish-NLP/t5-small-nl24-casing-punctuation-correction'    
tokenizer_t5 = AutoTokenizer.from_pretrained(model_checkpoint)
model_t5 = export_and_get_onnx_model(model_checkpoint)
#model_t5 = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint, from_flax=False, torch_dtype=torch.float32).to(device)

    
    

videos_out_path = Path("./videos_out")
videos_out_path.mkdir(parents=True, exist_ok=True)

samples_data = sorted(Path('examples').glob('*.json'))
SAMPLES = []
for file in samples_data:
    with open(file) as f:
        sample = json.load(f)
    SAMPLES.append(sample)
VIDEOS = list(map(lambda x: [x['video']], SAMPLES))

total_inferences_since_reboot = 0
total_cuts_since_reboot = 0




async def speech_to_text(video_file_path):
    """
    Takes a video path to convert to audio, transcribe audio channel to text  timestamps

    Using https://huggingface.co/tasks/automatic-speech-recognition pipeline
    """
    global total_inferences_since_reboot
    if(video_file_path == None):
        raise ValueError("Error no video input")

    video_path = Path(video_file_path)
    
    try:
        # convert video to audio 16k using PIPE to audio_memory
        audio_memory, _ = ffmpeg.input(video_path).output(
            '-', format="wav", ac=1, ar='16k').overwrite_output().global_args('-loglevel', 'quiet').run(capture_stdout=True)
    except Exception as e:
        raise RuntimeError("Error converting video to audio")

    last_time = time.time()

    try:
        
        output = speech_recognizer(
            audio_memory, return_timestamps="word",  chunk_length_s=10, stride_length_s=(4, 2))
        
        transcription = output["text"].lower()
        
        timestamps = [[chunk["text"].lower(), chunk["timestamp"][0], chunk["timestamp"][1]]
                        for chunk in output['chunks']]
        input_ids = tokenizer_t5(transcription, return_tensors="pt").input_ids.to(device)
        outputs = model_t5.generate(input_ids, max_length=128)
        case_corrected_text = tokenizer_t5.decode(outputs[0], skip_special_tokens=True)
        translated = model.generate(**tokenizer_marian([case_corrected_text], return_tensors="pt", padding=True))
        translated_plain = "".join([tokenizer_marian.decode(t, skip_special_tokens=True) for t in translated])

        for timestamp in timestamps:
            total_inferences_since_reboot += 1
        
    
        df = pd.DataFrame(timestamps, columns = ['word', 'start','stop'])
    
        df['start'] = df['start'].astype('float16')
        df['stop'] = df['stop'].astype('float16')
        

        print("\n\ntotal_inferences_since_reboot: ",
                total_inferences_since_reboot, "\n\n")
        return (transcription, transcription, timestamps,df, case_corrected_text, translated_plain)
    except Exception as e:
        raise RuntimeError("Error Running inference with local model", e)


def create_srt(text_out_t5, df):
    
    df.columns = ['word', 'start', 'stop']
    
    df_sentences = pd.DataFrame(columns=['sentence','start','stop','translated'])
    found_match_value = 0
    found_match_word = ""
    
    t5_sentences = re.split('[.]|[?]|[!]', text_out_t5)
    t5_sentences = [sentence.replace('.','').replace('?','').replace('!','') for sentence in t5_sentences if sentence]
    
    for i, sentence in enumerate(t5_sentences):
        sentence = sentence.lower().split(" ")
        if i == 0:
            df_subset = df[df['stop'] <10]
            start = df.iloc[0]['start']
        
            for j, word in enumerate(df_subset['word']):
              temp_value = fuzz.partial_ratio((word), sentence[-1])
              if temp_value > found_match_value:
                found_match_value = temp_value
                found_match_word = word
        
            stop = df_subset[df_subset['word'] == found_match_word]
            
            translated = model.generate(**tokenizer_marian(t5_sentences[i], return_tensors="pt", padding=True))
            translated_plain = [tokenizer_marian.decode(t, skip_special_tokens=True) for t in translated]
            
            dict_to_add = {
                'sentence': t5_sentences[i],
                'start': start,
                'stop': stop.iloc[0]['stop'],
                'translated': translated_plain[0]
            }
            
            df_sentences = df_sentences.append(dict_to_add, ignore_index=True)
            new_start = df.iloc[stop.index.values[0]+1]['start']
            new_stop = new_start + 10
        else:
            found_match_value = 0
            found_match_word = ""
            
            df_subset = df[(df['start'] >= new_start) & (df['stop'] <= new_stop)]
            start = df_subset.iloc[0]['start']
            
            for j, word in enumerate(df_subset['word']):
                temp_value = fuzz.partial_ratio((word), sentence[-1])
                if temp_value > found_match_value:
                    found_match_value = temp_value
                    found_match_word = word
            stop = df_subset[df_subset['word'] == found_match_word]
            
                
            translated = model.generate(**tokenizer_marian(t5_sentences[i], return_tensors="pt", padding=True))
            translated_plain = [tokenizer_marian.decode(t, skip_special_tokens=True) for t in translated]
            
            
            dict_to_add = {
                'sentence': t5_sentences[i],
                'start': start,
                'stop': stop.iloc[0]['stop'],
                'translated': translated_plain[0]
            }
            df_sentences = df_sentences.append(dict_to_add, ignore_index=True)
            try:
                new_start = df.iloc[stop.index.values[0]+1]['start']
                new_stop = new_start + 10
            except Exception as e:
                df_sentences = df_sentences.iloc[0:i+1]
    
    return df_sentences

def create_srt_and_burn(video_in, srt_sentences):
    srt_sentences.columns = ['sentence', 'start', 'stop','translated']
    srt_sentences.dropna(inplace=True)
    srt_sentences['start'] = srt_sentences['start'].astype('float')
    srt_sentences['stop'] = srt_sentences['stop'].astype('float')
    
    
    with open('testi.srt','w') as file:
        for i in range(len(srt_sentences)):
            file.write(str(i+1))
            file.write('\n')
            start = (time.strftime('%H:%M:%S', time.gmtime(srt_sentences.iloc[i]['start'])))
            if "." in str(srt_sentences.iloc[i]['start']):
                if len(str(srt_sentences.iloc[i]['start']).split('.')[1]) > 3:
                    start = start + '.' + str(srt_sentences.iloc[i]['start']).split('.')[1][:3]
                else:
                    start = start + '.' + str(srt_sentences.iloc[i]['start']).split('.')[1]
            file.write(start)
            stop = (time.strftime('%H:%M:%S', time.gmtime(srt_sentences.iloc[i]['stop'])))
            if len(str(srt_sentences.iloc[i]['stop']).split('.')[1]) > 3:
                stop = stop + '.' + str(srt_sentences.iloc[i]['stop']).split('.')[1][:3]    
            else:
                stop = stop + '.' + str(srt_sentences.iloc[i]['stop']).split('.')[1]
            file.write(' --> ')
            file.write(stop)
            file.write('\n')
            file.writelines(srt_sentences.iloc[i]['translated'])
            if int(i) != len(srt_sentences)-1:
                file.write('\n\n')
    try:
        file1 = open('./testi.srt', 'r')
        Lines = file1.readlines()
        
        count = 0
        # Strips the newline character
        for line in Lines:
            count += 1
                
        
        
        video_out = str(Path(video_in)).replace('.mp4', '_out.mp4')
        command = "ffmpeg -i {} -y -vf subtitles=./testi.srt {}".format(Path(video_in), Path(video_out))
        os.system(command)
        return video_out
    except Exception as e:
        print(e)
        return video_out


# ---- Gradio Layout -----
video_in = gr.Video(label="Video file", interactive=True)
text_in = gr.Textbox(label="Transcription", lines=10, interactive=True)
text_out_t5 = gr.Textbox(label="Transcription T5", lines=10, interactive=True)
translation_out = gr.Textbox(label="Translation", lines=10, interactive=True)
text_out_timestamps = gr.Textbox(label="Word level timestamps", lines=10, interactive=True)
srt_sentences = gr.DataFrame(label="Srt lines", row_count=(0, "dynamic"))
video_out = gr.Video(label="Video Out")
diff_out = gr.HighlightedText(label="Cuts Diffs", combine_adjacent=True)
examples = gr.components.Dataset(
    components=[video_in], samples=VIDEOS, type="index")

demo = gr.Blocks(enable_queue=True, css='''
#cut_btn, #reset_btn { align-self:stretch; }
#\\31 3 { max-width: 540px; }
.output-markdown {max-width: 65ch !important;}
''')
demo.encrypt = False
with demo:
    transcription_var = gr.Variable()
    timestamps_var = gr.Variable()
    timestamps_df = gr.Dataframe(visible=False, row_count=(0, "dynamic"))
    with gr.Row():
        with gr.Column():
            gr.Markdown('''
            # Create videos with English subtitles from videos spoken in Finnish
            This project is a quick proof of concept of a simple video editor where you can add English subtitles to Finnish videos.
            This space currently only works for short videos (Up to 128 tokens) but will be improved in next versions.
            Space uses our finetuned Finnish ASR models, Our pretrained + finetuned Finnish T5 model for casing+punctuation correction and Opus-MT models from Helsinki University for Finnish --> English translation.
            This space was inspired by https://huggingface.co/spaces/radames/edit-video-by-editing-text
            ''')

    with gr.Row():

        examples.render()

        def load_example(id):
            video = SAMPLES[id]['video']
            transcription = ''
            timestamps = SAMPLES[id]['timestamps']

            return (video, transcription, transcription, timestamps)

        examples.click(
            load_example,
            inputs=[examples],
            outputs=[video_in, text_in, transcription_var, timestamps_var],
            queue=False)
    with gr.Row():
        with gr.Column():
            video_in.render()
            transcribe_btn = gr.Button("1. Press here to transcribe Audio")
            transcribe_btn.click(speech_to_text, [video_in], [
                text_in, transcription_var, text_out_timestamps,timestamps_df, text_out_t5, translation_out])

    with gr.Row():
        gr.Markdown('''
        ### Here you will get varying outputs from different parts of the processing
        ASR model output, T5 model output which corrects casing + hyphenation, sentence level translations and word level timestamps''')

    with gr.Row():
        with gr.Column():
            text_in.render()
        with gr.Column():
            text_out_t5.render()
        with gr.Column():
            translation_out.render()
        with gr.Column():
            text_out_timestamps.render()
    with gr.Row():
        with gr.Column():
            translate_and_make_srt_btn = gr.Button("2. Press here to create rows for subtitles")
            translate_and_make_srt_btn.click(create_srt, [text_out_t5, timestamps_df], [
                srt_sentences])
    with gr.Row():
        with gr.Column():
            srt_sentences.render()
    with gr.Row():
        with gr.Column():
            translate_and_make_srt_btn = gr.Button("3. Press here to create subtitle file and insert translations to video")
            translate_and_make_srt_btn.click(create_srt_and_burn, [video_in, srt_sentences], [
                video_out])
            video_out.render()
                
if __name__ == "__main__":
    demo.launch(debug=True)