createvideo / utils.py
meraj12's picture
Create utils.py
f998b37 verified
import whisper
import os
import torch
import clip
from PIL import Image
from moviepy.editor import *
from transformers import pipeline
from langdetect import detect
from gtts import gTTS
import tempfile
device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("ViT-B/32", device=device)
whisper_model = whisper.load_model("small")
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en")
def transcribe_audio_segments(audio_path):
result = whisper_model.transcribe(audio_path, word_timestamps=True)
segments = result["segments"]
return segments
def translate_text(text, target_lang="en"):
detected = detect(text)
if detected != target_lang:
translated = translator(text)[0]['translation_text']
return translated
return text
def match_images(text, image_paths):
text = clip.tokenize([text]).to(device)
images = [preprocess(Image.open(img_path)).unsqueeze(0).to(device) for img_path in image_paths]
images_tensor = torch.cat(images)
with torch.no_grad():
text_features = clip_model.encode_text(text)
image_features = clip_model.encode_image(images_tensor)
similarity = (image_features @ text_features.T).squeeze()
best_index = similarity.argmax().item()
return image_paths[best_index]
def generate_speech(text, lang="en"):
tts = gTTS(text, lang=lang)
temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
tts.save(temp_audio.name)
return temp_audio.name
def create_video_segments(segments, audio_path, image_paths, output_path="final_video.mp4"):
clips = []
for seg in segments:
segment_text = seg['text'].strip()
translated = translate_text(segment_text)
duration = seg['end'] - seg['start']
matched_img = match_images(translated, image_paths)
image_clip = ImageClip(matched_img).set_duration(duration)
txt_clip = TextClip(translated, fontsize=30, color='white', bg_color='black', size=image_clip.size)
txt_clip = txt_clip.set_duration(duration).set_position(("center", "bottom"))
audio = AudioFileClip(audio_path).subclip(seg['start'], seg['end'])
final_clip = CompositeVideoClip([image_clip, txt_clip])
final_clip = final_clip.set_audio(audio)
clips.append(final_clip)
final_video = concatenate_videoclips(clips, method="compose")
final_video.write_videofile(output_path, fps=24)
return output_path