|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, WhisperFeatureExtractor
|
|
from moviepy.editor import VideoFileClip, AudioFileClip
|
|
import nltk
|
|
nltk.download('punkt', quiet=True)
|
|
from nltk.tokenize import sent_tokenize
|
|
|
|
|
|
def transcribe(video_file, transcribe_to_text=True, transcribe_to_srt=True, target_language='en'):
|
|
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-large-v3"
|
|
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)
|
|
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_id)
|
|
|
|
video = VideoFileClip(video_file)
|
|
audio = video.audio
|
|
duration = audio.duration
|
|
chunk_duration = 60
|
|
n_chunks = int(np.ceil(duration / chunk_duration))
|
|
|
|
full_transcription = ""
|
|
for i in range(n_chunks):
|
|
start_time = i * chunk_duration
|
|
end_time = min((i + 1) * chunk_duration, duration)
|
|
|
|
audio_chunk = audio.subclip(start_time, end_time)
|
|
temp_file_path = f"temp_audio_chunk_{i}.wav"
|
|
audio_chunk.write_audiofile(temp_file_path, codec='pcm_s16le')
|
|
|
|
sound_array = AudioFileClip(temp_file_path).to_soundarray(fps=16000)
|
|
if sound_array.ndim > 1:
|
|
sound_array = np.mean(sound_array, axis=1)
|
|
|
|
input_features = feature_extractor(sound_array, sampling_rate=16000, return_tensors="pt").input_features
|
|
input_features = input_features.to(device=device, dtype=torch_dtype)
|
|
|
|
with torch.no_grad():
|
|
if target_language:
|
|
model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language=target_language,
|
|
task="transcribe")
|
|
generated_ids = model.generate(input_features, max_length=448)
|
|
transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
|
|
|
full_transcription += transcription + " "
|
|
os.remove(temp_file_path)
|
|
print(f"Processed chunk {i + 1}/{n_chunks}")
|
|
|
|
|
|
sentences = sent_tokenize(full_transcription.strip())
|
|
|
|
|
|
total_chars = sum(len(s) for s in sentences)
|
|
sentence_times = []
|
|
current_time = 0
|
|
for sentence in sentences:
|
|
sentence_duration = (len(sentence) / total_chars) * duration
|
|
sentence_times.append((current_time, current_time + sentence_duration))
|
|
current_time += sentence_duration
|
|
|
|
output = ""
|
|
if transcribe_to_text:
|
|
output += "Text Transcription:\n" + full_transcription + "\n\n"
|
|
|
|
if transcribe_to_srt:
|
|
output += "SRT Transcription:\n"
|
|
for i, (sentence, (start, end)) in enumerate(zip(sentences, sentence_times), 1):
|
|
output += f"{i}\n{format_time(start)} --> {format_time(end)}\n{sentence}\n\n"
|
|
|
|
return output
|
|
|
|
def format_time(seconds):
|
|
m, s = divmod(seconds, 60)
|
|
h, m = divmod(m, 60)
|
|
return f"{int(h):02d}:{int(m):02d}:{s:06.3f}".replace('.', ',') |
