Add YouTube download fallback and direct file upload option

app.py

@@ -41,14 +41,15 @@ def extract_audio(url):
             os.rename(out_file, audio_file)
             return audio_file, video_title, temp_dir
         else:
-            st.error("No audio stream found")
-            return None, None, None
+            st.warning("No audio stream found with pytube. Trying yt-dlp...")
+            return extract_audio_with_ytdlp(url)
     except Exception as e:
         st.warning(f"pytube error: {e}. Falling back to yt-dlp.")
         return extract_audio_with_ytdlp(url)
 
 def extract_audio_with_ytdlp(url):
     try:
+        temp_dir = tempfile.mkdtemp()
         ydl_opts = {
             'format': 'bestaudio/best',
             'postprocessors': [{
@@ -57,15 +58,54 @@ def extract_audio_with_ytdlp(url):
                 'preferredquality': '192',
             }],
             'outtmpl': '%(title)s.%(ext)s',
-            'paths': {'home': tempfile.mkdtemp()}
+            'paths': {'home': temp_dir},
+            # Add options to help bypass YouTube restrictions
+            'nocheckcertificate': True,
+            'ignoreerrors': True,
+            'no_warnings': True,
+            'quiet': True,
+            'cookiefile': None,  # Use cookies if available
+            'extractor_args': {'youtube': {'player_client': ['web']}},  # Use web player
+            # Add more random headers to avoid detection
+            'http_headers': {
+                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.82 Safari/537.36',
+                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
+                'Accept-Language': 'en-US,en;q=0.5',
+                'Accept-Encoding': 'gzip, deflate',
+                'DNT': '1',
+                'Connection': 'keep-alive',
+            }
         }
         with yt_dlp.YoutubeDL(ydl_opts) as ydl:
             info_dict = ydl.extract_info(url, download=True)
+            if not info_dict:
+                st.error("Could not process this YouTube URL. The video might be age-restricted or protected.")
+                return None, None, None
+                
             video_title = info_dict.get('title', None)
-            audio_file = os.path.join(ydl_opts['paths']['home'], f"{video_title}.mp3")
-            return audio_file, video_title, ydl_opts['paths']['home']
+            if not video_title:
+                st.error("Could not retrieve video title")
+                return None, None, None
+                
+            audio_file = os.path.join(temp_dir, f"{video_title}.mp3")
+            if not os.path.exists(audio_file):
+                # Try alternate filename format which might be used by yt-dlp
+                possible_files = [f for f in os.listdir(temp_dir) if f.endswith('.mp3')]
+                if possible_files:
+                    audio_file = os.path.join(temp_dir, possible_files[0])
+                else:
+                    st.error("Downloaded file not found")
+                    return None, None, None
+                    
+            return audio_file, video_title, temp_dir
     except Exception as e:
-        st.error(f"An error occurred with yt-dlp: {e}")
+        error_message = str(e)
+        if "Sign in to confirm you're not a bot" in error_message:
+            st.error("YouTube is detecting automated access. Try a different video or try again later.")
+        elif "This video is unavailable" in error_message:
+            st.error("This video is unavailable or may be restricted.")
+        else:
+            st.error(f"An error occurred with yt-dlp: {e}")
         return None, None, None
 
 
@@ -352,12 +392,15 @@ def make_predictions(model, processed_audio, audio_features, url, video_name):
     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)]
 
-    st.write(f"**Video Title:** {video_name}")
+    st.write(f"**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 {format_time(start_time)} to {format_time(end_time)}: [{link}]({link})")
+    for i, (start_time, end_time) in enumerate(zip(chorus_start_times, chorus_end_times)):
+        if url:  # Only create YouTube timestamp links if we have a URL
+            link = f"{url}&t={int(start_time)}s"
+            st.write(f"Chorus {i+1}: {format_time(start_time)} to {format_time(end_time)}: [{link}]({link})")
+        else:
+            st.write(f"Chorus {i+1}: {format_time(start_time)} to {format_time(end_time)}")
 
     if len(chorus_start_times) == 0:
         st.write("No choruses identified.")
@@ -419,26 +462,59 @@ def main():
     st.image(cover_image, use_column_width=True)
     st.title("Chorus Finder")
     st.write("This app uses a pre-trained convolutional recurrent neural network to predict chorus locations in music. To learn more about this project, visit [github.com/dennisvdang/chorus-detection](https://github.com/dennisvdang/chorus-detection).")
-    st.write("Enter a YouTube song URL to find the chorus!")
-    url = st.text_input("YouTube URL")
-    if st.button("Find Chorus"):
-        if url:
-            with st.spinner('Analyzing YouTube link...'):
-                audio_file, video_title, temp_dir = extract_audio(url)
-            if audio_file:
-                with st.spinner('Analyzing YouTube link...'):
-                    strip_silence(audio_file)
-                with st.spinner('Processing audio...'):
-                    processed_audio, audio_features = process_audio(audio_path=audio_file)
-                with st.spinner('Loading model...'):
-                    model = load_model(MODEL_PATH)
-                with st.spinner('Making predictions...'):
-                    smoothed_predictions = make_predictions(model, processed_audio, audio_features, url, video_title)
-                with st.spinner('Plotting predictions...'):
-                    plot_predictions(audio_features, smoothed_predictions)
-                shutil.rmtree(temp_dir)
+    
+    st.write("### Option 1: YouTube URL")
+    url = st.text_input("Enter a YouTube URL")
+    process_youtube = st.button("Find Chorus from YouTube")
+    
+    st.write("### Option 2: Upload Audio File")
+    st.write("If YouTube downloading fails, you can upload an MP3 file directly")
+    uploaded_file = st.file_uploader("Choose an MP3 file", type=['mp3'])
+    process_upload = st.button("Find Chorus from Upload")
+    
+    if process_youtube and url:
+        with st.spinner('Analyzing YouTube link...'):
+            audio_file, video_title, temp_dir = extract_audio(url)
+        if audio_file:
+            with st.spinner('Processing audio...'):
+                # Ensure we strip silence for YouTube downloads
+                strip_silence(audio_file)
+                processed_audio, audio_features = process_audio(audio_path=audio_file, trim_silence=False)
+            with st.spinner('Loading model...'):
+                model = load_model(MODEL_PATH)
+            with st.spinner('Making predictions...'):
+                smoothed_predictions = make_predictions(model, processed_audio, audio_features, url, video_title)
+            with st.spinner('Plotting predictions...'):
+                plot_predictions(audio_features, smoothed_predictions)
+            shutil.rmtree(temp_dir)
         else:
+            st.error("Failed to download from YouTube. Please try uploading the audio file directly instead.")
+    
+    elif process_upload and uploaded_file is not None:
+        # Save the uploaded file to a temporary directory
+        with st.spinner('Processing uploaded file...'):
+            temp_dir = tempfile.mkdtemp()
+            temp_path = os.path.join(temp_dir, "uploaded_audio.mp3")
+            with open(temp_path, "wb") as f:
+                f.write(uploaded_file.getbuffer())
+            
+            # Use the filename as the video title
+            file_title = uploaded_file.name.replace('.mp3', '')
+            
+            # Process the uploaded file
+            strip_silence(temp_path)
+            processed_audio, audio_features = process_audio(audio_path=temp_path, trim_silence=False)
+            model = load_model(MODEL_PATH)
+            # Since we don't have a YouTube URL for the uploaded file, pass an empty string
+            smoothed_predictions = make_predictions(model, processed_audio, audio_features, "", file_title)
+            plot_predictions(audio_features, smoothed_predictions)
+            shutil.rmtree(temp_dir)
+    
+    elif (process_youtube and not url) or (process_upload and uploaded_file is None):
+        if process_youtube and not url:
             st.error("Please enter a valid YouTube URL")
+        if process_upload and uploaded_file is None:
+            st.error("Please upload an MP3 file")
 
 if __name__ == "__main__":
     main()