sortify/sortify.py at master · wardzin/sortify · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import spotipy
import spotipy.util
import numpy as np
import pandas as pd
from tqdm.autonotebook import tqdm, trange


client_id = 'b69a9985fa8842deb0691b2d0e3f0b69'
client_secret = 'd9e9ae2924174c139a5a9ccb303f9f3a'
redirect_uri = 'http://localhost/'

username = '22mrmbu7oumkrb56tcsclawdi' # Daniel

token = spotipy.util.prompt_for_user_token(username, 'user-library-read playlist-modify-private', client_id=client_id, client_secret=client_secret, redirect_uri=redirect_uri)
sp = spotipy.Spotify(auth=token)


def chunks(n, x):
    return [x[i:i + n] for i in range(0, len(x), n)]


def get_user_tracks():
    tracks = []

    offset = 0
    pbar = tqdm()
    while True:
        track_set = sp.current_user_saved_tracks(limit=50, offset=offset)
        if track_set:
            track_set = track_set['items']
            pbar.update(len(track_set))
            tracks += track_set
            if len(track_set) < 50:
                pbar.close()
                break
        else:
            pbar.close()
            break
        offset += 50

    track_data = {
        'id': [],
        'name': [],
        'artist': [],
    }
    for track in tracks:
        if 'track' not in track:
            continue
        track = track['track']
        track_data['id'].append(track['id'])
        track_data['name'].append(track['name'])
        if 'remix' in track['name'].lower():
            track_data['artist'].append(track['artists'][-1]['id'])
        else:
            track_data['artist'].append(track['artists'][0]['id'])

    return pd.DataFrame(track_data).set_index('id')


def get_track_features(tracks):
    feature_data = {
        'id': [],
        'danceability': [],
        'energy': [],
        'key': [],
        'loudness': [],
        'mode': [],
        'speechiness': [],
        'acousticness': [],
        'instrumentalness': [],
        'liveness': [],
        'valence': [],
        'tempo': [],
        'time_signature': [],
    }

    for chunk in tqdm(chunks(50, tracks)):
        track_set = sp.audio_features(chunk)
        for track_features in track_set:
            for key in feature_data.keys():
                feature_data[key].append(track_features[key])

    return pd.DataFrame(feature_data).set_index('id')


def get_timbre_data(tracks):
    timbre_data = []

    for track in tqdm(tracks):
        timbres = []
        durations = []
        for segment in sp.audio_analysis(track)['segments']:
            timbres.append(segment['timbre'])
            durations.append(segment['duration'])
        timbre_data.append(np.average(timbres, axis=0, weights=durations))

    return pd.DataFrame(timbre_data, index=tracks)


def get_artist_data(artists):
    artist_data = {
        'id': [],
        'name': [],
        'popularity': [],
        'genres': []
    }

    for chunk in tqdm(chunks(50, artists)):
        for artist in sp.artists(chunk)['artists']:
            artist_data['id'].append(artist['id'])
            artist_data['name'].append(artist['name'])
            artist_data['popularity'].append(artist['popularity'])
            artist_data['genres'].append(artist['genres'])

    return pd.DataFrame(artist_data).set_index('id')


def get_related_artists(artists):
    related_artists = {}

    for artist in tqdm(artists):
        related_artist_list = sp.artist_related_artists(artist)['artists']

        if len(related_artist_list) <= 1:
            continue

        related_artists[artist] = []
        for related_artist in related_artist_list:
            related_artists[artist].append(related_artist['id'])

    return related_artists