main.py

import os
import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
from pydub import AudioSegment
import soundfile as sf

from ai_dev.engine import run_diarization
from ai_dev.utils import waveplot, combined_waveplot, rename

temp_path = "./tempDir"

st.set_page_config(
    page_title = 'Gnosis',
    page_icon = ":crocodile:",
    layout = "wide",
    menu_items = {
        'Get Help': 'https://www.google.com/',
        'Report a bug': 'https://github.com/GnosisUM/Challenge-2_BuildWithAI/issues/new',
        'About': 
            "Submission by Team Gnosis for #BuildWithAI 2021 Hackathon Challenge 2"
    },
)

m = st.markdown("""
<style>
div.stButton > button:first-child {
    background-color: rgb(204, 49, 49);
}
</style>""", unsafe_allow_html=True)

# function to convert .mp3 file to .wav file
def mp3_to_wav(mp3_file_name: str, path: str):
    sound = AudioSegment.from_file(path+'/'+mp3_file_name, format='mp3')
    sound.export(path+'/'+mp3_file_name[:-4]+'.wav', format="wav")
    os.remove(path+'/'+mp3_file_name)

# function to convert .mp4 file to .wav file
def mp4_to_wav(mp4_file_name: str, path: str):
    sound = AudioSegment.from_file(path+'/'+mp4_file_name, format='mp4')
    sound.export(path+'/'+mp4_file_name[:-4]+'.wav', format="wav")
    os.remove(path+'/'+mp4_file_name)

# function to compute audio using AI model
def compute_audio(wav_file):
    csv_file = None # placeholder for model output
    return csv_file

# function to get filename from a list
def get_file_names(files):
    file_names = []
    for i in range(len(files)):
        file_names.append(files[i].name)
    return file_names

# function to display audio player based on dropdown selection
def display_audio_playback(file_name, files_list):
    for i in range(len(files_list)):
        if file_name == files_list[i].name:
            st.audio(files_list[i])

# function to save uploaded file(s) into a temporary directory
def save_upload(file):
    with open(os.path.join("tempDir", file.name),"wb") as f:
         f.write(file.getbuffer())

import random
def generate_attributes():

    v_attr1 = ["Masculine", "Feminine"]
    for i in range(2):
        st.write(f"**Speaker {i}'s Attributes**")
        attr1 = {
            'Voice Attributes': [random.choice(v_attr1)],
            'Confidence': [random.uniform(0.5, 1)]
        }
        df = pd.DataFrame(attr1, index=None)
        # df.index.name = 'Voice Attributes'
        st.table(df)

#function to display attributes generated by AI model
def display_attributes(WAV_FILE):
    try:
        with st.spinner('Loading attributes...'):
            df, segments = run_diarization(WAV_FILE=WAV_FILE, embed_model='xvec', cluster_method='sc')

        df.drop(['start_sample', 'end_sample'], axis=1, inplace=True)
        df = df.rename(columns={"start":"Start (seconds)", "end":"End (seconds)", "label":"Speaker No."})
        # df.index.name = 'Speaker No.'
        # df.set_index(['Speaker No.'], inplace=True)
        st.table(df)
        signal, sr = sf.read(WAV_FILE)
        generate_attributes()
        #st.pyplot(waveplot(signal, sr))
        st.pyplot(combined_waveplot(signal, sr, segments))
        
    except RuntimeError:
        print("Something went wrong!")
        try:
            os.listdir(temp_path)
        except FileNotFoundError:
            st.error("File not found! Try reuploading the file.")

st.image('./resource/Vokal.-colorized.png', width=275)
st.write("A program powered by A.I. model to detect the number of speakers and generate attributes related to the audio uploaded.")

# create a temporary directory if not existed
if not os.path.isdir(temp_path):
    os.mkdir(temp_path)
    
col1, col2 = st.columns(2)

with col1:
    # File uploader
    uploaded_files = st.file_uploader(
        "Upload audio file(s)",
        type=['wav','mp3','mp4'],

        accept_multiple_files=True,
        help="Choose to upload single or multiple files of format .wav, .mp3 and .mp4"
    )

    # Save uploaded files into the created temporary directory 
    for i in range(len(uploaded_files)):
        uploaded_files[i].name = rename(uploaded_files[i].name)
        save_upload(uploaded_files[i])

    # If the audio file is an mp3 or mp4, convert to .wav for AI model
    temp_list = get_file_names(uploaded_files)
    for i in range(len(temp_list)):
        if temp_list[i].endswith('.mp3'):
            mp3_to_wav(temp_list[i], temp_path)
        elif temp_list[i].endswith('.mp4'):
            mp4_to_wav(temp_list[i], temp_path)
    
    st.write('**Remember to clear uploaded files before exiting the program.**')

    # Clear tempDir folder and its content
    if st.button("Clear uploaded files"):
        for f in os.listdir(temp_path):
            os.remove(os.path.join(temp_path, f))
        os.rmdir(temp_path)

        st.error("Audio files deleted!")

with col2:
    # Displays dropdown menu if number of files > 1
    if len(uploaded_files) > 1:
        dropdown_selection = st.selectbox(
            "Choose audio file to play",
            get_file_names(uploaded_files),
            1,
            help="Choose any uploaded file(s) for playback and to see attributes generated by our A.I."
        )
        
        WAV_FILE = os.path.join(temp_path, dropdown_selection)
        display_audio_playback(dropdown_selection, uploaded_files)
        with st.expander('See attributes'):
            display_attributes(WAV_FILE)
            # plot_waveform(dropdown_selection, uploaded_files)

    # Displays only the audio player when number of files == 1
    if len(uploaded_files) == 1:
        st.audio(uploaded_files[0])
        WAV_FILE = os.path.join(temp_path, temp_list[0])
        with st.expander('See attributes'):
            display_attributes(WAV_FILE)

# for f in os.listdir(temp_path):
#     os.remove(os.path.join(temp_path, f))
# os.rmdir(temp_path)