Upload 4 files

.streamlit/config.toml

@@ -0,0 +1,10 @@
+[theme]
+base = "dark"
+primaryColor = "#FFFFFF"
+backgroundColor = "#212121"
+secondaryBackgroundColor = "#757575"
+textColor = "#FFFFFF"
+font = "sans serif"
+
+[browser]
+gatherUsageStats = false

ui.py

@@ -0,0 +1,97 @@
+import json
+
+import streamlit as st
+from ui_backend import (
+    check_for_api,
+    cut_audio_file,
+    display_predictions,
+    load_audio,
+    predict_multiple,
+    predict_single,
+)
+
+
+def main():
+    # Page settings
+    st.set_page_config(
+        page_title="Music Instrument Recognition", page_icon="🎸", layout="wide", initial_sidebar_state="collapsed"
+    )
+
+    # Sidebar
+    with st.sidebar:
+        st.title("⚙️ Settings")
+        selected_model = st.selectbox(
+            "Select Model",
+            ("Accuracy", "Speed"),
+            index=0,
+            help="Select a slower but more accurate model or a faster but less accurate model",
+        )
+
+    # Main title
+    st.markdown(
+        "<h1 style='text-align: center; color: #FFFFFF; font-size: 3rem;'>Instrument Recognition 🎶</h1>",
+        unsafe_allow_html=True,
+    )
+
+    # Upload widget
+    audio_file = load_audio()
+
+    # Send a health check request to the API in a loop until it is running
+    api_running = check_for_api(10)
+
+    # Enable or disable a button based on API status
+    predict_valid = False
+    cut_valid = False
+
+    if api_running:
+        st.info("API is running", icon="🤖")
+
+    if audio_file:
+        num_files = len(audio_file)
+        st.write(f"Number of uploaded files: {num_files}")
+        predict_valid = True
+        if len(audio_file) > 1:
+            cut_valid = False
+        else:
+            audio_file = audio_file[0]
+            cut_valid = True
+            name = audio_file.name
+
+    if cut_valid:
+        cut_audio = st.checkbox(
+            "✂️ Cut duration",
+            disabled=not predict_valid,
+            help="Cut a long audio file. Model works best if audio is around 15 seconds",
+        )
+
+        if cut_audio:
+            audio_file = cut_audio_file(audio_file, name)
+
+    result = st.button("Predict", disabled=not predict_valid, help="Send the audio to API to get a prediction")
+
+    if result:
+        predictions = {}
+        if isinstance(audio_file, list):
+            predictions = predict_multiple(audio_file, selected_model)
+
+        else:
+            predictions = predict_single(audio_file, name, selected_model)
+
+        # Sort the dictionary alphabetically by key
+        sorted_predictions = dict(sorted(predictions.items()))
+
+        # Convert the sorted dictionary to a JSON string
+        json_string = json.dumps(sorted_predictions)
+        st.download_button(
+            label="Download JSON",
+            file_name="predictions.json",
+            mime="application/json",
+            data=json_string,
+            help="Download the predictions in JSON format",
+        )
+
+        display_predictions(sorted_predictions)
+
+
+if __name__ == "__main__":
+    main()

ui_backend.py

@@ -0,0 +1,254 @@
+import io
+import os
+import time
+from json import JSONDecodeError
+import math
+
+import requests
+import soundfile as sf
+import streamlit as st
+
+if os.environ.get("IS_DOCKER", False):
+    backend = "http://api:7860"
+else:
+    backend = "http://0.0.0.0:7860"
+
+INSTRUMENTS = {
+    "tru": "Trumpet",
+    "sax": "Saxophone",
+    "vio": "Violin",
+    "gac": "Acoustic Guitar",
+    "org": "Organ",
+    "cla": "Clarinet",
+    "flu": "Flute",
+    "voi": "Voice",
+    "gel": "Electric Guitar",
+    "cel": "Cello",
+    "pia": "Piano",
+}
+
+
+def load_audio():
+    """
+    Upload a WAV audio file and display it in a Streamlit app.
+
+    :return: A BytesIO object representing the uploaded audio file, or None if no file was uploaded.
+    :rtype: Optional[BytesIO]
+    """
+
+    audio_file = st.file_uploader(label="Upload audio file", type="wav", accept_multiple_files=True)
+    if len(audio_file) > 0:
+        st.audio(audio_file[0])
+        return audio_file
+    else:
+        return None
+
+
+@st.cache_data(show_spinner=False)
+def check_for_api(max_tries: int):
+    """
+    Check if the API is running by making a health check request.
+
+    :param max_tries: The maximum number of attempts to check the API's health.
+    :type max_tries: int
+    :return: True if the API is running, False otherwise.
+    :rtype: bool
+    """
+    trial_count = 0
+
+    with st.spinner("Waiting for API..."):
+        while trial_count <= max_tries:
+            try:
+                response = health_check()
+                if response:
+                    return True
+            except requests.exceptions.ConnectionError:
+                trial_count += 1
+                # Handle connection error, e.g. API not yet running
+                time.sleep(5)  # Sleep for 1 second before retrying
+        st.error("API is not running. Please refresh the page to try again.", icon="🚨")
+        st.stop()
+
+
+def cut_audio_file(audio_file, name):
+    """
+    Cut an audio file and return the cut audio data as a tuple.
+
+    :param audio_file: The path of the audio file to be cut.
+    :type audio_file: str
+    :param name: The name of the audio file to be cut.
+    :type name: str
+    :raises RuntimeError: If the audio file cannot be read.
+    :return: A tuple containing the name and the cut audio data as a BytesIO object.
+    :rtype: tuple
+    """
+    try:
+        audio_data, sample_rate = sf.read(audio_file)
+    except RuntimeError as e:
+        raise e
+
+    # Display audio duration
+    duration = round(len(audio_data) / sample_rate, 2)
+    st.info(f"Audio Duration: {duration} seconds")
+
+    # Get start and end time for cutting
+    start_time = st.number_input("Start Time (seconds)", min_value=0.0, max_value=duration - 1, step=0.1)
+    end_time = st.number_input("End Time (seconds)", min_value=start_time, value=duration, max_value=duration, step=0.1)
+
+    # Convert start and end time to sample indices
+    start_sample = int(start_time * sample_rate)
+    end_sample = int(end_time * sample_rate)
+
+    # Cut audio
+    cut_audio_data = audio_data[start_sample:end_sample]
+
+    # Create a temporary in-memory file for cut audio
+    audio_file = io.BytesIO()
+    sf.write(audio_file, cut_audio_data, sample_rate, format="wav")
+
+    # Display cut audio
+    st.audio(audio_file, format="audio/wav")
+    audio_file = (name, audio_file)
+
+    return audio_file
+
+
+def display_predictions(predictions: dict):
+    """
+    Display the predictions using instrument names instead of codes.
+
+    :param predictions: A dictionary containing the filenames and instruments detected in them.
+    :type predictions: dict
+    """
+
+    # Display the results using instrument names instead of codes
+    for filename, instruments in predictions.items():
+        st.subheader(filename)
+
+        if isinstance(instruments, str):
+            st.write(instruments)
+
+        else:
+            with st.container():
+                col1, col2 = st.columns([1, 3])
+                present_instruments = [
+                    INSTRUMENTS[instrument_code] for instrument_code, presence in instruments.items() if presence
+                ]
+                if present_instruments:
+                    for instrument_name in present_instruments:
+                        with col1:
+                            st.write(instrument_name)
+                        with col2:
+                            st.write("✔️")
+                else:
+                    st.write("No instruments found in this file.")
+
+
+def health_check():
+    """
+    Sends a health check request to the API and checks if it's running.
+
+    :return: Returns True if the API is running, else False.
+    :rtype: bool
+    """
+
+    # Send a health check request to the API
+    response = requests.get(f"{backend}/health-check", timeout=100)
+
+    # Check if the API is running
+    if response.status_code == 200:
+        return True
+    else:
+        return False
+
+
+def predict(data, model_name):
+    """
+    Sends a POST request to the API with the provided data and model name.
+
+    :param data: The audio data to be used for prediction.
+    :type data: bytes
+    :param model_name: The name of the model to be used for prediction.
+    :type model_name: str
+    :return: The response from the API.
+    :rtype: requests.Response
+    """
+
+    file = {"file": data}
+    request_data = {"model_name": model_name}
+
+    response = requests.post(
+        f"{backend}/predict", params=request_data, files=file, timeout=300
+    )  # Replace with your API endpoint URL
+
+    return response
+
+
+@st.cache_data(show_spinner=False)
+def predict_single(audio_file, name, selected_model):
+    """
+    Predicts the instruments in a single audio file using the selected model.
+
+    :param audio_file: The audio file to be used for prediction.
+    :type audio_file: bytes
+    :param name: The name of the audio file.
+    :type name: str
+    :param selected_model: The name of the selected model.
+    :type selected_model: str
+    :return: A dictionary containing the predicted instruments for the audio file.
+    :rtype: dict
+    """
+
+    predictions = {}
+
+    with st.spinner("Predicting instruments..."):
+        response = predict(audio_file, selected_model)
+
+        if response.status_code == 200:
+            prediction = response.json()["prediction"]
+            predictions[name] = prediction.get(name, "Error making prediction")
+        else:
+            st.write(response)
+            try:
+                st.json(response.json())
+            except JSONDecodeError:
+                st.error(response.text)
+            st.stop()
+    return predictions
+
+
+@st.cache_data(show_spinner=False)
+def predict_multiple(audio_files, selected_model):
+    """
+    Generates predictions for multiple audio files using the selected model.
+
+    :param audio_files: A list of audio files to make predictions on.
+    :type audio_files: List[UploadedFile]
+    :param selected_model: The model to use for making predictions.
+    :type selected_model: str
+    :return: A dictionary where the keys are the names of the audio files and the values are the predicted labels.
+    :rtype: Dict[str, str]
+    """
+
+    predictions = {}
+    progress_text = "Getting predictions for all files. Please wait."
+    progress_bar = st.empty()
+    progress_bar.progress(0, text=progress_text)
+
+    num_files = len(audio_files)
+
+    for i, file in enumerate(audio_files):
+        name = file.name
+        response = predict(file, selected_model)
+        if response.status_code == 200:
+            prediction = response.json()["prediction"]
+            predictions[name] = prediction[name]
+            progress_bar.progress((i + 1) / num_files, text=progress_text)
+        else:
+            predictions[name] = "Error making prediction."
+    progress_bar.empty()
+    return predictions
+
+
+if __name__ == "__main__":
+    pass