script for downloading MTG-QBH dataset collection

getSongFromYoutube.py

import os
import click
import re
import csv
import numpy as np
from selenium import webdriver
from bs4 import BeautifulSoup
from urllib.parse import quote_plus

YOUTUBE_QUERY = 'https://www.youtube.com/results?search_query='
EXT = 'ogg'
FILE_FORMAT = 'vorbis'


def search_videos(query, driver):
    while True:
        driver.get(YOUTUBE_QUERY + query)
        driver.implicitly_wait(1.2)
        try:
            elements = driver.find_elements_by_xpath('//*[@id="video-title"]')
            links = []
            for element in elements:
                title = element.get_attribute('title')
                href = element.get_attribute('href')
                if href is not None:
                    links.append((title, href))
            assert len(links) > 0
            break
        except:
            _ = input('[PROMPT] Continue?')
            print(elements)
    return links


def autoSelection(song, links):
    def jaccard_similarity(list1, list2):
        intersection = len(list(set(list1).intersection(list2)))
        union = (len(list1) + len(list2)) - intersection
        return float(intersection / union)
    ref = song[0].lower().split() + song[1].lower().split() + ['-']
    score = [jaccard_similarity(ref, link[0].lower().strip().split())
             for link in links]
    if score[0] > 0.35:
        # prefer first result
        idx = 0
    else:
        idx = np.argmax(np.array(score))
    if score[idx] < 0.8:
        print(f'[WARNING] low jaccard similarity:{score[idx]}')
    return idx


def interactiveSelection(song, links):
    print('Found:\n'+'\n'.join(
        [f'[{i}]: {link[0].strip()}' for i, link in enumerate(links)]))
    choice = input('Pick one: ')
    while not(choice.isdigit()) or not(0 <= int(choice) < len(links)):
        print("Oups, that was wrong. Try again!")
        choice = input('Pick one: ')
    return int(choice)


def writeURL(output, songs, auto=False):
    with open(output, newline='') as csvfile:
        # <artist>\t<title>\t<label>\t<url>\t<video_title>
        reader = csv.reader(csvfile, delimiter=',', quotechar='\"')
        completed = [row[:3] for row in reader][1:]
    songs = list(filter(lambda x: x not in completed, songs))
    print('[INFO] starting webdriver')
    driver = webdriver.Chrome()
    driver.get('https://www.youtube.com')
    print('[INFO] webdriver initialized')

    queries = [quote_plus(' - '.join(song[:2]).strip()) for song in songs]
    with open(output, newline='', mode='a') as csvfile:
        writer = csv.writer(csvfile, delimiter=',',
                            quotechar='\"', quoting=csv.QUOTE_MINIMAL)
        try:
            for song, query in zip(songs, queries):
                print(f'[SEARCH]: {song} @ {query}')
                links = search_videos(query, driver)
                idx = autoSelection(
                    song, links) if auto else interactiveSelection(song, links)
                video_title, url = links[idx]
                print(f'[SELECTED]: {video_title}')
                writer.writerow(list(song) + [url, video_title])
        except KeyboardInterrupt:
            pass


def downloadSong(artist, title, label, url, v_title, dst):
    target = os.path.join(dst, f'{artist} - {title}.webm')
    command = f'youtube-dl --extract-audio --audio-format {FILE_FORMAT} --output \"{target}\" \'{url}\' '
    print(f'[INFO] Downloading {v_title} @ {url}')
    os.system(command)


@click.group()
def cli():
    pass


@click.command()
@click.argument('songurls', nargs=1, type=click.Path(exists=True))
@click.argument('dst', nargs=1, type=click.Path(exists=True))
def download(songurls, dst):
    with open(songurls, newline='') as csvfile:
        # <artist>\t<title>\t<label>\t<url>\t<video_title>
        reader = csv.reader(csvfile, delimiter=',', quotechar='\"')
        records = [row for row in reader][1:]
    completed = os.listdir(dst)
    for record in records:
        if f'{record[0]} - {record[1]}.{EXT}' not in completed:
            downloadSong(*record, dst)
        else:
            print(f'[INFO] skipping {record[-1]}')
    for record in records:
        if record[-1]+f'.{EXT}' in completed:
            # <TO DO ?>add metadata
            pass


@click.command()
@click.argument('songlist', nargs=1, type=click.Path(exists=True))
@click.option('--auto', default=False, type=click.BOOL, help='match urls auto selection')
@click.option('--output', default='./songUrls.csv', type=click.Path())
def geturl(songlist, auto, output):
    with open(songlist, newline='') as csvfile:
        # <artist>\t<title>\t<label>
        reader = csv.reader(csvfile, delimiter=',', quotechar='\"')
        songs = [row for row in reader][1:]
    if not os.path.exists(output):
        with open(output, newline='', mode='a') as csvfile:
            # <artist>\t<title>\t<label>\t<url>\t<video_title>
            writer = csv.writer(csvfile, delimiter=',',
                                quotechar='\"', quoting=csv.QUOTE_MINIMAL)
            writer.writerow(['artist', 'title', 'label', 'url', 'video_title'])
    writeURL(output, songs, auto)


cli.add_command(download)
cli.add_command(geturl)
if __name__ == '__main__':
    cli()

songUrls.csv