Spaces:
Sleeping
Sleeping
import gradio as gr | |
from pytube import YouTube | |
from moviepy.editor import VideoFileClip | |
import torch | |
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline | |
device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
model_id = "openai/whisper-tiny" | |
model = AutoModelForSpeechSeq2Seq.from_pretrained( | |
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True | |
) | |
model.to(device) | |
processor = AutoProcessor.from_pretrained(model_id) | |
pipe = pipeline( | |
"automatic-speech-recognition", | |
model=model, | |
tokenizer=processor.tokenizer, | |
feature_extractor=processor.feature_extractor, | |
max_new_tokens=128, | |
chunk_length_s=15, | |
batch_size=16, | |
return_timestamps=True, | |
torch_dtype=torch_dtype, | |
device=device, | |
) | |
# get video and extract video | |
def get_video(yt_url): | |
try: | |
video = YouTube(yt_url) | |
video.streams.get_by_itag(22).download(filename='video.mp4') | |
print('Video succesfully downloaded from Youtube') | |
except Exception as e: | |
print(f'Failed to download Youtube video \nerror : {e}') | |
def audio_from_video(video_path): | |
try: | |
video = VideoFileClip(video_path) | |
audio = video.audio | |
audio.write_audiofile('audio.wav') | |
video.close() | |
audio.close() | |
except Exception as e: | |
print(f'Failed to extract audio from video \nerror : {e}') | |
def transcribe(url): | |
video_path = './video.mp4' | |
get_video(url) | |
audio_from_video(video_path) | |
# transcribe audio | |
audio = 'audio.wav' | |
text_audio = pipe(audio) | |
chunks = text_audio['chunks'] | |
chunks_count = len(chunks) | |
chunk_id = [] | |
timestamps = [] | |
texts = [] | |
start_time = [] | |
end_time = [] | |
for i in range(0, chunks_count): | |
chunk_id.append(i) | |
texts.append(chunks[i]['text']) | |
start_time.append(chunks[i]['timestamp'][0]) | |
end_time.append(chunks[i]['timestamp'][1]) | |
chunk_length = [] | |
for i in range(0, chunks_count-1): | |
chunk_length.append(round(end_time[i] - start_time[i], 3)) | |
output = list(zip(chunk_id, chunk_length, texts, start_time, end_time)) | |
sample_output_list = [] | |
for sublist in output: | |
chunk_dict = { | |
"chunk_id": sublist[0], | |
"chunk_length": sublist[1], | |
"text": sublist[2], | |
"start_time": sublist[3], | |
"end_time": sublist[4] | |
} | |
sample_output_list.append(chunk_dict) | |
return sample_output_list | |
intf = gr.Interface( | |
fn=transcribe, | |
inputs = ["text"], | |
outputs = ["text"] | |
) | |
intf.launch(inline=False) |