CS370-Milestone5 / Milestone5API.py
JaMugen's picture
Update Milestone5API.py
0d79804
Raw
History Blame
2.6 kB
import os
import re
from pytube import YouTube
from youtube_transcript_api import YouTubeTranscriptApi
from transformers import MarianMTModel, MarianTokenizer
def download_video_transcript(video_url, source_lang='en', target_lang='fr'):
try:
video_id = re.search(r"(?<=v=)[\w-]+", video_url)
if video_id:
video_id = video_id.group()
yt = YouTube(video_url)
stream = yt.streams.get_highest_resolution()
print(f'Downloading {yt.title}...')
modified_title = yt.title.replace(" ", "_")
download_path = 'CS370_Milestone5/videos' # Updated path to CS370_Milestone5 repository
captions_path = 'CS370_Milestone5/captions' # Updated path to CS370_Milestone5 repository
# Create directories if they don't exist
os.makedirs(download_path, exist_ok=True)
os.makedirs(captions_path, exist_ok=True)
video_file = f'{download_path}/{modified_title}.mp4'
stream.download(output_path=download_path, filename=modified_title + '.mp4')
print('Download completed!')
transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
transcript = transcript_list.find_generated_transcript([source_lang]).fetch()
original_captions = ""
for i, line in enumerate(transcript):
start_time = line['start']
formatted_time = f"{int(start_time // 60):02d}:{int(start_time % 60):02d}"
original_captions += f"{formatted_time} {line['text']}\n"
original_filename = f'{captions_path}/{modified_title}_original.txt'
with open(original_filename, 'w', encoding='utf-8') as file:
file.write(original_captions)
# Translation part
translated_captions = translate_text_file(original_filename,
f'{captions_path}/{modified_title}_translated.txt',
source_lang=source_lang,
target_lang=target_lang)
# Delete video file after captions are saved
os.remove(video_file)
return original_captions, translated_captions, original_filename, f'{captions_path}/{modified_title}_translated.txt'
else:
print("Video ID not found in URL.")
return None, None, None, None
except Exception as e:
print("Error:", e)
return None, None, None, None