Update app.py

app.py

@@ -119,18 +119,23 @@ def extract_audio_features(data_tracks, legit_genres):
     all_tracks_uris = all_tracks_uris[valid_indexes]
     all_audio_features = np.array(all_audio_features)[valid_indexes]
     all_tracks_audio_features = dict(zip(relevant_audio_features, [[audio_f[k] for audio_f in all_audio_features] for k in relevant_audio_features]))
-    genres = dict()
+    all_tracks_genres = []
+    indexes_by_genre = dict()
     for index, uri in enumerate(all_tracks_uris):
         track = data_tracks[uri]
         track_genres = track['track']['genres']
-        for g in track_genres:
-            if g not in genres.keys():
-                genres[g] = [index]
+        all_tracks_genres.append([])
+        for glabel in track_genres:
+            legit_genre = find_legit_genre(glabel, legit_genres)
+            if legit_genre in indexes_by_genre.keys():
+                indexes_by_genre[legit_genre].append(index)
             else:
-                genres[g].append(index)
-    genres = aggregate_genres(genres, legit_genres)
-    genres_labels = sorted(genres.keys())
-    return all_tracks_uris, all_tracks_audio_features, genres, genres_labels
+                indexes_by_genre[legit_genre] = [index]
+            all_tracks_genres[-1].append(legit_genre)
+        all_tracks_genres[-1] = sorted(set(all_tracks_genres[-1]))
+    genres_labels = sorted(indexes_by_genre.keys())
+    all_tracks_genres = np.array(all_tracks_genres)
+    return all_tracks_uris, all_tracks_audio_features, all_tracks_genres, indexes_by_genre, genres_labels
     # st.session_state['music_extracted'] = dict(all_tracks_uris=all_tracks_uris,
     #                                            all_tracks_audio_features=all_tracks_audio_features,
     #                                            genres=genres,
@@ -149,11 +154,15 @@ def select_songs(legit_genres):
     users_playlists = "Add a list of user urls, one per line (optional)"
     users_links = st.text_area(users_playlists, value="")
     label_playlists = "Add a list of playlists urls, one per line (optional)"
-    playlist_links = st.text_area(label_playlists, value="https://open.spotify.com/playlist/1H7a4q8JZArMQiidRy6qon")
+    playlist_links = st.text_area(label_playlists, value="https://open.spotify.com/playlist/1H7a4q8JZArMQiidRy6qon\nhttps://open.spotify.com/playlist/6wbaZqht4w6CMv3od5taax?si=5c6ebe13fdd049b6")
     extract_button = centered_button(st.button, 'Extract music', n_columns=5)
 
-    all_tracks_uris, all_tracks_audio_features, genres, genres_labels = [None] * 4
+    all_tracks_uris, all_tracks_audio_features, all_tracks_genres, indexes_by_genre, genres_labels = [None] * 5
     if extract_button or debug or 'extract_button' in st.session_state.keys():
+        if extract_button:
+            updated_sources = True
+        else:
+            updated_sources = False
         st.session_state['extract_button'] = True
         # check the user input music sources
         if playlist_links == "" and users_links == "":
@@ -183,13 +192,13 @@ def select_songs(legit_genres):
             if len(data_tracks.keys()) < 10:
                 st.warning('Please select more music sources.')
             else:
-                all_tracks_uris, all_tracks_audio_features, genres, genres_labels = extract_audio_features(data_tracks, legit_genres)
+                all_tracks_uris, all_tracks_audio_features, all_tracks_genres, indexes_by_genre, genres_labels = extract_audio_features(data_tracks, legit_genres)
                 print(f'4. audio feature extraction: {time.time() - init_time:.2f}')
                 print(f'\t total extraction: {time.time() - init_time_tot:.2f}')
                 st.success(f'{len(data_tracks.keys())} tracks found!')
-    return all_tracks_uris, all_tracks_audio_features, genres, genres_labels
+    return all_tracks_uris, all_tracks_audio_features, all_tracks_genres, indexes_by_genre, genres_labels, updated_sources
 
-def customize_widgets(genres_labels):
+def customize_widgets(genres_labels, updated_sources):
     st.subheader("Step 3: Customize it!")
     st.markdown('##### Which genres?')
 
@@ -203,8 +212,9 @@ def customize_widgets(genres_labels):
         with columns[3]:
             uncheck_all = st.button('Uncheck all')
 
-        if 'checkboxes' not in st.session_state.keys():
+        if 'checkboxes' not in st.session_state.keys() or updated_sources:
             st.session_state['checkboxes'] = [True] * len(genres_labels)
+            updated_sources = False
 
         empty_checkboxes = wall_of_checkboxes(genres_labels, max_width=5)
         if check_all:
@@ -228,12 +238,12 @@ def customize_widgets(genres_labels):
     return target_mood, exploration
 
 @st.cache
-def filter_songs_by_genre(checkboxes, genres_labels, genres):
+def filter_songs_by_genre(checkboxes, genres_labels, indexes_by_genre):
     # filter songs by genres
     selected_labels = [genres_labels[i] for i in range(len(genres_labels)) if checkboxes[i]]
     genre_selected_indexes = []
     for label in selected_labels:
-        genre_selected_indexes += genres[label]
+        genre_selected_indexes += indexes_by_genre[label]
     genre_selected_indexes = np.array(sorted(set(genre_selected_indexes)))
     return genre_selected_indexes
 
@@ -246,7 +256,7 @@ def find_best_songs_for_mood(all_tracks_audio_features, genre_selected_indexes,
     return min_dist_indexes, n_candidates
 
 @st.cache
-def run_exploration(selected_tracks_uris, playlist_length, exploration, all_tracks_uris, target_mood, reauthenticate):
+def run_exploration(selected_tracks_uris, selected_tracks_genres, playlist_length, exploration, all_tracks_uris, target_mood, reauthenticate):
     # sample exploration songs
     if exploration > 0:
         n_known = int(playlist_length * (1 - exploration))
@@ -254,6 +264,7 @@ def run_exploration(selected_tracks_uris, playlist_length, exploration, all_trac
         print(f'Number of new songs: {n_new}, known songs: {n_known}')
         known_songs = selected_tracks_uris[:n_known]
         seed_songs = selected_tracks_uris[-n_new:]
+        seed_genres = selected_tracks_genres[-n_new:]
         dict_args = dict()  # enforce bounds on recommendations' moods
         for i_m, m in enumerate(['valence', 'energy', 'danceability']):
             dict_args[f'min_{m}'] = max(0, target_mood[i_m] - 0.1)
@@ -268,18 +279,20 @@ def run_exploration(selected_tracks_uris, playlist_length, exploration, all_trac
             dict_args_loose[f'max_{m}'] = min(1, target_mood[i_m] + 0.3)
         new_songs = []
         counter_seed = 0
-        counter_failure = 0
         while len(new_songs) < n_new:
             try:
                 print(seed_songs[counter_seed])
                 print(dict_args)
-                reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], market="from_token", **dict_args)['tracks']
+                reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], seed_genres=seed_genres[counter_seed],
+                                          market="from_token", country='from_token', **dict_args)['tracks']
                 if len(reco) == 0:
                     print('Using loose bounds')
-                    reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], market="from_token", **dict_args_loose)['tracks']
+                    reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], seed_genres=seed_genres[counter_seed],
+                                              market="from_token", country='from_token', **dict_args_loose)['tracks']
                     if len(reco) == 0:
                         print('Using looser bounds')
-                        reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], market="from_token", **dict_args_looser)['tracks']
+                        reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], seed_genres=seed_genres[counter_seed],
+                                                  market="from_token", country='from_token', **dict_args_looser)['tracks']
                         if len(reco) == 0:
                             print('Removing bounds')
                             reco = sp.recommendations(seed_tracks=[seed_songs[counter_seed]], market="from_token")['tracks']
@@ -300,17 +313,20 @@ def run_exploration(selected_tracks_uris, playlist_length, exploration, all_trac
     return selected_tracks_uris
 
 @st.cache
-def sample_playlist(n_candidates, playlist_length, genre_selected_indexes, min_dist_indexes, all_tracks_uris):
+def sample_playlist(n_candidates, playlist_length, genre_selected_indexes, min_dist_indexes, all_tracks_uris, all_tracks_genres):
     # give more freedom to randomize the playlist
     if n_candidates > 5 * playlist_length:
         selected_tracks_indexes = genre_selected_indexes[min_dist_indexes[:int(playlist_length * 2)]]
 
     else:
         selected_tracks_indexes = genre_selected_indexes[min_dist_indexes[:playlist_length]]
-    selected_tracks_uris = all_tracks_uris[selected_tracks_indexes]
-    np.random.shuffle(selected_tracks_uris)
+    shuffled_indexes = np.arange(len(selected_tracks_indexes))
+    np.random.shuffle(shuffled_indexes)
+    selected_tracks_uris = all_tracks_uris[selected_tracks_indexes][shuffled_indexes]
+    selected_tracks_genres = all_tracks_genres[selected_tracks_indexes][shuffled_indexes]
     selected_tracks_uris = selected_tracks_uris[:playlist_length]
-    return selected_tracks_uris
+    selected_tracks_genres = selected_tracks_genres[:playlist_length]
+    return selected_tracks_uris, selected_tracks_genres
 
 def run_app():
     global sp
@@ -333,16 +349,16 @@ def run_app():
         return sp
 
     if 'login' in st.session_state or debug:
-        all_tracks_uris, all_tracks_audio_features, genres, genres_labels = select_songs(legit_genres)
+        all_tracks_uris, all_tracks_audio_features, all_tracks_genres, indexes_by_genre, genres_labels, updated_sources = select_songs(legit_genres)
 
         if all_tracks_uris is not None:
-            target_mood, exploration = customize_widgets(genres_labels)
+            target_mood, exploration = customize_widgets(genres_labels, updated_sources)
             custom_button = centered_button(st.button, 'Run customization', n_columns=5)
-            if custom_button or 'run_custom' in st.session_state.keys():
+            if custom_button or 'run_custom' in st.session_state.keys() or debug:
                 st.session_state['run_custom'] = True
                 checkboxes = st.session_state['checkboxes'].copy()
                 init_time = time.time()
-                genre_selected_indexes = filter_songs_by_genre(checkboxes, genres_labels, genres)
+                genre_selected_indexes = filter_songs_by_genre(checkboxes, genres_labels, indexes_by_genre)
                 if len(genre_selected_indexes) < 10:
                     genre_selected_indexes = None
                     st.warning('Please select more genres or add more music sources.')
@@ -361,7 +377,8 @@ def run_app():
                         playlist_length = st.number_input(f'Pick a playlist length, given {n_candidates} candidates.', min_value=5,
                                                           value=min(10, n_candidates//3), max_value=n_candidates//3)
 
-                        selected_tracks_uris = sample_playlist(n_candidates, playlist_length, genre_selected_indexes, min_dist_indexes, all_tracks_uris)
+                        selected_tracks_uris, selected_tracks_genres = sample_playlist(n_candidates, playlist_length, genre_selected_indexes,
+                                                                                       min_dist_indexes, all_tracks_uris, all_tracks_genres)
                         print(f'8. Sample songs: {time.time() - init_time:.2f}')
                         init_time = time.time()
 
@@ -371,7 +388,8 @@ def run_app():
                         else:
                             generation_button = centered_button(st.button, 'Generate playlist', n_columns=5)
                             if generation_button:
-                                selected_tracks_uris = run_exploration(selected_tracks_uris, playlist_length, exploration, all_tracks_uris, target_mood.flatten(), reauthenticate)
+                                selected_tracks_uris = run_exploration(selected_tracks_uris, selected_tracks_genres, playlist_length, exploration, all_tracks_uris,
+                                                                       target_mood.flatten(), reauthenticate)
                                 print(f'9. run exploration: {time.time() - init_time:.2f}')
                                 init_time = time.time()