|
# 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')) |
i don't really know what is up with it, but when it creates all the files, it removes the last symbol, so when it tries to find the file for frame extraction it fails to find it.


as you can see on the image, the names for thumbnails are missing the last '6', so i had to manually add it to the mkv file, only to run into multiprocessing issue.
that's numbers 2, numbers 1 works fine (well, until a pickle)