Script fixes

app.py

@@ -18,6 +18,7 @@ from pytube import YouTube
 from sklearn.preprocessing import StandardScaler
 import shutil
 import streamlit as st
+import tempfile
 
 
 # Constants
@@ -29,26 +30,27 @@ N_FEATURES = 15
 MODEL_PATH = "models/CRNN/best_model_V3.h5"
 AUDIO_TEMP_PATH = "output/temp"
 
-def extract_audio(url, output_path=AUDIO_TEMP_PATH):
+
+def extract_audio(url):
     try:
         yt = YouTube(url)
         video_title = yt.title
         audio_stream = yt.streams.filter(only_audio=True).first()
         if audio_stream:
-            os.makedirs(output_path, exist_ok=True)
-            out_file = audio_stream.download(output_path)
+            temp_dir = tempfile.mkdtemp()
+            out_file = audio_stream.download(temp_dir)
             base, _ = os.path.splitext(out_file)
             audio_file = base + '.mp3'
             if os.path.exists(audio_file):
                 os.remove(audio_file)
             os.rename(out_file, audio_file)
-            return audio_file, video_title
+            return audio_file, video_title, temp_dir
         else:
             st.error("No audio stream found")
-            return None, None
+            return None, None, None
     except Exception as e:
         st.error(f"An error occurred: {e}")
-        return None, None
+        return None, None, None
 
 
 def strip_silence(audio_path):
@@ -412,38 +414,24 @@ def make_predictions(model, processed_audio, audio_features, url, video_name):
         audio_features.meter_grid, sr=audio_features.sr, hop_length=audio_features.hop_length)
     chorus_start_times = [meter_grid_times[i] for i in range(len(
         smoothed_predictions)) if smoothed_predictions[i] == 1 and (i == 0 or smoothed_predictions[i - 1] == 0)]
+    chorus_end_times = [meter_grid_times[i + 1] for i in range(len(
+        smoothed_predictions)) if smoothed_predictions[i] == 1 and (i == len(smoothed_predictions) - 1 or smoothed_predictions[i + 1] == 0)]
 
-    youtube_links = [
-        f"\033]8;;{url}&t={int(start_time)}s\033\\{url}&t={int(start_time)}s\033]8;;\033\\" for start_time in chorus_start_times
-    ]
-    max_length = max([len(link) for link in youtube_links] + [len(video_name), len(
-        f"Number of choruses identified: {len(chorus_start_times)}")] if chorus_start_times else [0])
-    header_footer = "=" * (max_length + 4)
-    print()
-    print()
-    print(header_footer)
-    print(f"{video_name.center(max_length + 2)}")
-    print(f"Number of choruses identified: {len(chorus_start_times)}".center(
-        max_length + 4))
-    print(header_footer)
-    for link in youtube_links:
-        print(link)
-    print(header_footer)
+    st.write(f"**Video Title:** {video_name}")
+    st.write(f"**Number of choruses identified:** {len(chorus_start_times)}")
+
+    for start_time, end_time in zip(chorus_start_times, chorus_end_times):
+        link = f"{url}&t={int(start_time)}s"
+        st.write(f"Chorus from {start_time:.2f}s to {end_time:.2f}s: [Link]({link})")
 
     if len(chorus_start_times) == 0:
-        print("No choruses identified.")
+        st.write("No choruses identified.")
 
     return smoothed_predictions
 
 
 def plot_meter_lines(ax: plt.Axes, meter_grid_times: np.ndarray) -> None:
-    """
-    Draw meter grid lines on the plot.
-
-    Parameters:
-    - ax (plt.Axes): The matplotlib axes object to draw on.
-    - meter_grid_times (np.ndarray): Array of times at which to draw the meter lines.
-    """
+    """Draw meter grid lines on the plot."""
     for time in meter_grid_times:
         ax.axvline(x=time, color='grey', linestyle='--',
                    linewidth=1, alpha=0.6)
@@ -499,17 +487,16 @@ def main():
     url = st.text_input("YouTube URL")
     if st.button("Find Chorus"):
         if url:
-            audio_file, video_title = extract_audio(url)
+            audio_file, video_title, temp_dir = extract_audio(url)
             if audio_file:
                 strip_silence(audio_file)
-                processed_audio, audio_features = process_audio(audio_path=AUDIO_TEMP_PATH)
+                processed_audio, audio_features = process_audio(audio_path=audio_file)
                 model = load_model()
                 smoothed_predictions = make_predictions(model, processed_audio, audio_features, url, video_title)
                 plot_predictions(audio_features=audio_features, predictions=smoothed_predictions)
-                shutil.rmtree(AUDIO_TEMP_PATH)
+                shutil.rmtree(temp_dir)
         else:
             st.error("Please enter a valid YouTube URL")
 
 if __name__ == "__main__":
     main()
-