Extract frames from youtube videos

fftube.py

import argparse
import youtube_dl
import ffmpy

samples_per_video = 10
game_dic = {}

def save_frame(url, time, game_name):
    if game_name not in game_dic.keys():
        game_dic[game_name] = 0
    output_file = game_name + '-' + str(game_dic[game_name]) + '.jpg'
    ff = ffmpy.FFmpeg(
        inputs={url: '-ss ' + str(time) + ' -t 1'},
        outputs={output_file: '-f mjpeg'})
    ff.run()
    game_dic[game_name] += 1

def grab_random_frames(youtube_link, game_name):
    ydl = youtube_dl.YoutubeDL({'outtmpl': '%(id)s%(ext)s'})
    with ydl:
        result = ydl.extract_info(youtube_link, download=False)
    video_link = result['formats'][-1]['url']
    duration = result['duration']

    for i in range(1, duration, int(duration/samples_per_video)):
        save_frame(video_link, i, game_name)


parser = argparse.ArgumentParser(description='download pile of videos from youtube')
parser.add_argument('--file', dest='file', action='store', required=True)
args = parser.parse_args()

with open(args.file) as f:
    content = f.readlines()

for line in content:
    url, game_name = str(line).strip().split(',')
    grab_random_frames(url, game_name)

frames-from-games.py

import argparse
import youtube_dl
import ffmpy
import os

sample_rate = 5
game_dic = {}
video_per_game = 3

def save_frame(url, time, game):
    output_file = os.path.join(game, game + '-' + str(game_dic[game]) + '.jpg')
    ff = ffmpy.FFmpeg(
        inputs={url: '-ss ' + str(time) + ' -t 1'},
        outputs={output_file: '-f mjpeg'})
    ff.run()
    game_dic[game] += 1

def search_and_extract(query, game_name, rate):
    print('Processing links for: {}'.format(game_name))
    ydl = youtube_dl.YoutubeDL({'outtmpl': '%(id)s%(ext)s'})
    with ydl:
        result = ydl.extract_info('ytsearch' + str(video_per_game) + ':' + str(query), download=False)
    
    if not os.path.exists(game_name):
        os.makedirs(game_name)

    for game_video in result['entries']:
        video_link = game_video['formats'][-1]['url']
        duration = game_video['duration']

        for i in range(1, duration, int(duration/rate)):
            save_frame(video_link, i, game_name)


parser = argparse.ArgumentParser(description='download pile of videos from youtube')
parser.add_argument('--file', dest='file', action='store', required=True)
args = parser.parse_args()

with open(args.file) as f:
    content = f.readlines()

for line in content:
    game_name, search_query = str(line).strip().split(', ')

    if game_name not in game_dic.keys():
        game_dic[game_name] = 0

    if search_query is not "":
        search_and_extract(search_query, game_name, sample_rate)