|
# main.py - Gets the video, exports the frames, and does some processing on |
|
# those frames. Also gets subtitles. |
|
|
|
# Check requirements.txt to install the things required for this script. |
|
# Also needs ffmpeg and imagemagik |
|
|
|
import youtube_dl |
|
import ffmpeg |
|
import glob |
|
import os |
|
from multiprocessing import Process |
|
|
|
|
|
def download_yt_video(url: str, options: dict): |
|
with youtube_dl.YoutubeDL(options) as ydl: |
|
ydl.download([url]) |
|
|
|
|
|
def not_exist_create(dir: str): |
|
if not os.path.exists(dir): |
|
print(f'{dir} does not exist. Will create') |
|
os.mkdir(dir) |
|
|
|
def stream_export_raw_frames(in_file: str, out_file: str): |
|
ffmpeg.input(in_file).output(out_file).run() |
|
|
|
def stream_export_gamma_frames(in_file: str, out_file: str): |
|
ffmpeg.input(in_file).filter('eq', gamma=3, saturation=1.3).output(out_file).run() |
|
|
|
def export_ffmpeg_frames(video_id: str, in_dir: str, out_dir: str, |
|
out_dir_gamma: str): |
|
print('Extracting frames...') |
|
f = os.path.join(in_dir, f'{video_id}.mkv') |
|
if not os.path.exists(f): |
|
print('cannot find downloaded file') |
|
return |
|
|
|
not_exist_create(out_dir) |
|
not_exist_create(out_dir_gamma) |
|
|
|
raw_template_path = os.path.join(out_dir, '%03d.bmp') |
|
gamma_template_path = os.path.join(out_dir_gamma, '%03d.bmp') |
|
|
|
proc_raw = Process(target=stream_export_raw_frames, args=(f, raw_template_path)) |
|
proc_gamma = Process(target=stream_export_gamma_frames, args=(f, gamma_template_path)) |
|
|
|
proc_raw.start() |
|
proc_gamma.start() |
|
|
|
proc_raw.join() |
|
proc_gamma.join() |
|
|
|
|
|
if __name__ == '__main__': |
|
video_id = 'ymYFqNUt05g' |
|
|
|
# leave these alone |
|
out_dir = 'out' |
|
url = f'https://www.youtube.com/watch?v={video_id}' |
|
opts = { |
|
'writedescription': True, |
|
'writeinfojson': True, |
|
'writesubtitles': True, |
|
'write_all_thumbnails': True, |
|
'allasubtitles': True, |
|
'keepvideo': True, |
|
'keepaudio': True, |
|
'format': 'bestvideo+bestaudio', |
|
'outtmpl': os.path.join(out_dir, '%(id)s.%(ext)s'), |
|
'merge_output_format': 'mkv' |
|
} |
|
download_yt_video(url, opts) |
|
export_ffmpeg_frames(video_id, out_dir, os.path.join(out_dir, 'frames'), |
|
os.path.join(out_dir, 'gamma_frames')) |
Thanks 👍