Updated to use model

app.py

@@ -14,6 +14,10 @@ import os
 import shutil
 import pandas as pd
 import plotly.express as px
+import torch
+from pyannote.audio import Pipeline
+from pyannote.core import Annotation, Segment
+from df.enhance import enhance, init_df
 
 PARQUET_DATASET_DIR = Path("parquet_dataset")
 PARQUET_DATASET_DIR.mkdir(parents=True,exist_ok=True)
@@ -23,6 +27,18 @@ sample_data = [f"CHEM1402_gt/24F_CHEM1402_Night_Class_Week_{i}_gt.rttm" for i in
 
 scheduler = ps.ParquetScheduler(repo_id="Sonogram/SampleDataset")
 
+secondDifference = config['secondDifference']
+gainWindow = config['gainWindow']
+minimumGain = config['minimumGain']
+maximumGain = config['maximumGain']
+attenLimDB = config['attenLimDB']
+
+# Instantiate and prepare model for training.
+dfModel, dfState, _ = init_df(model_base_dir="DeepFilterNet3")
+#dfModel.to(torch.device("cuda:0"))
+pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
+#pipeline.to(torch.device("cuda:1"))
+
 # Store results for viewing and further processing
 if 'results' not in st.session_state:
     st.session_state.results = []
@@ -52,7 +68,23 @@ def save_data(
     # Send to scheduler
     scheduler.append(data)
 
-st.set_page_config(layout="wide")
+def processFile(filePath):
+    global attenLimDb
+    global gainWindow
+    global minimumGain
+    global maximumGain
+    waveformList, sampleRate = su.splitIntoTimeSegments(filePath)
+    enhancedWaveformList = []
+    for w in waveformList:
+        newW = enhance(dfModel,dfState,w,atten_lim_db=attenLimDB)#.detach().cpu()
+        enhancedWaveformList.append(newW)
+    waveformEnhanced = combineWaveforms(enhancedWaveformList)     
+    waveform_gain_adjusted = equalizeVolume()(waveformEnhanced,sample_rate,gainWindow,minimumGain,maximumGain)
+    annotations = pipeline({"waveform": waveformEnhanced, "sample_rate": sample_rate})
+    speakerList = su.annotationToSpeakerList(annotations)
+    return speakerList, annotations
+
+#st.set_page_config(layout="wide")
 st.title("Lecturer Support Tool")
 
 uploaded_file_paths = st.file_uploader("Upload an audio of classroom activity to analyze", accept_multiple_files=True)
@@ -84,21 +116,39 @@ if uploaded_file_paths is not None:
     while (len(st.session_state.summaries) < len(valid_files)):
         st.session_state.summaries.append([])
     st.info(f'{len(valid_files)} valid files: {[fi.name for fi in valid_files]}')
+    
+if tab.button("Analyze All Audio",key=f"button_all"):
+    if len(valid_files) == 0:
+        tab.error('Upload file(s) first!')
+    else:
+        totalFiles = len(valid_files)
+        for i in range(totalFiles):
+            with st.spinner(text=f'Analyzing File {i+1} of {totalFiles}'):
+                # Text files use sample data
+                if file_paths[i].endswith('.txt'):
+                    time.sleep(1)
+                    # RTTM load as filler
+                    speakerList, annotations = su.loadAudioRTTM(sample_data[i])
+                    st.session_state.results[i] = (speakerList,annotations)
+                    st.session_state.summaries[i] = []
+                else:
+                    st.session_state.results[i] = processFile(file_paths[i])
+                    st.session_state.summaries[i] = []
+            
 for i, tab in enumerate(audio_tabs):
     if tab.button("Analyze Audio",key=f"button_{i}"):
-        if uploaded_file is None:
-            tab.error('Upload a file first!')
-        else:
-            # Process
-            # Pretend to take time as an example
+        # Text files use sample data
+        if file_paths[i].endswith('.txt'):
             with st.spinner(text='NOT ACTUALLY ANALYZING, JUST A FILLER ANIMATION'):
                 time.sleep(1)
-                st.success('Done')
-                
-            # RTTM load as filler
-            speakerList, annotations = su.loadAudioRTTM(sample_data[i])
-            st.session_state.results[i] = (speakerList,annotations)
-            st.session_state.summaries[i] = []
+                # RTTM load as filler
+                speakerList, annotations = su.loadAudioRTTM(sample_data[i])
+                st.session_state.results[i] = (speakerList,annotations)
+                st.session_state.summaries[i] = []
+        else:
+            with st.spinner(text='Analyzing File'):
+                st.session_state.results[i] = processFile(file_paths[i])
+                st.session_state.summaries[i] = []
             
     if len(st.session_state.results) > i and len(st.session_state.summaries) > i and len(st.session_state.results[i]) > 0:
         with st.spinner(text='Loading results...'):