diff --git a/summarize.py b/summarize.py index 867585c..e064231 100755 --- a/summarize.py +++ b/summarize.py @@ -1,13 +1,11 @@ #! /usr/bin/env python3 +import argparse import io import json import yt_dlp from ollama import chat, ChatResponse, Client -#video_url = 'https://youtu.be/jl4HOY8ZaEA' -video_url = 'https://www.youtube.com/watch?v=kTctVqjhDEw' - ydl_opts = { 'writesubtitles': True, # Enable downloading subtitles 'subtitleslangs': ['en'], # Specify subtitle language(s) @@ -69,33 +67,52 @@ def concatenate_subtitles(subtitle_json): # Join all collected text with spaces and return return ' '.join(result) -with yt_dlp.YoutubeDL(ydl_opts) as ydl: - # Extract metadata without downloading the video - info = ydl.extract_info(video_url, download=False) +def main(): + parser = argparse.ArgumentParser( + description="Download subtitles from a video and summarize it using a local Ollama instance." + ) + parser.add_argument( + 'url', + metavar='URL', + type=str, + help="The URL of the video to process." + ) - # Check if subtitles are available - subtitle_lang = 'en' # Change this to your desired language - subtitles_available = info.get('subtitles', {}) - automatic_subtitles_available = info.get('automatic_captions', {}) + # Parse out arguments + args = parser.parse_args() + video_url = args.url - if subtitle_lang in subtitles_available: - print(f"Downloading manual subtitles for language: {subtitle_lang}") - subtitle_url = subtitles_available[subtitle_lang][0]['url'] - elif subtitle_lang in automatic_subtitles_available: - print(f"No manual subtitles available. Falling back to auto-generated subtitles for language: {subtitle_lang}") - subtitle_url = automatic_subtitles_available[subtitle_lang][0]['url'] - else: - print(f"No subtitles (manual or auto-generated) available for language: {subtitle_lang}") - subtitle_url = None + with yt_dlp.YoutubeDL(ydl_opts) as ydl: + # Extract metadata without downloading the video + info = ydl.extract_info(video_url, download=False) + print(f"Summarizing video: {info.get('title', 'Unknown Title')}...") - # If a subtitle URL was found, download the subtitles - if subtitle_url: - subtitle_data = ydl.urlopen(subtitle_url).read().decode('utf-8') - else: - print("Failed to download subtitles.") - exit(50) + # Check if subtitles are available + subtitle_lang = 'en' # Change this to your desired language + subtitles_available = info.get('subtitles', {}) + automatic_subtitles_available = info.get('automatic_captions', {}) - subs = concatenate_subtitles(json.loads(subtitle_data)) - print("Getting summary...") - summary = get_summary(subs) - print(summary) + if subtitle_lang in subtitles_available: + print(f"Downloading manual subtitles for language: {subtitle_lang}...") + subtitle_url = subtitles_available[subtitle_lang][0]['url'] + elif subtitle_lang in automatic_subtitles_available: + print(f"No manual subtitles available. Falling back to auto-generated subtitles for language: {subtitle_lang}...") + subtitle_url = automatic_subtitles_available[subtitle_lang][0]['url'] + else: + print(f"No subtitles (manual or auto-generated) available for language: {subtitle_lang}!") + subtitle_url = None + exit(51) + + # If a subtitle URL was found, download the subtitles + if subtitle_url: + subtitle_data = ydl.urlopen(subtitle_url).read().decode('utf-8') + else: + print("Failed to download subtitles!") + exit(50) + + subs = concatenate_subtitles(json.loads(subtitle_data)) + print("Getting summary...") + summary = get_summary(subs) + print(summary) + +main()