Added loading of pretrained model

app.py

@@ -16,6 +16,10 @@ import streamlit as st
 import torch
 import pandas as pd
 from pyannote.audio import Pipeline
+from pyannote.audio.pipelines import SpeakerDiarization
+from pyannote.audio.models.segmentation import PyanNet
+from pyannote.audio import Inference
+from pyannote.pipeline.parameter import ParamDict
 
 import sonogram_utility as su
 import utils
@@ -71,7 +75,24 @@ if ENABLE_DENOISE:
 else:
     dfModel = dfState = None
 
+# Load baseline model
 pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
+
+# Generate replacement segmentation model
+newSpecs = pipeline._segmentation.model.specifications
+segModel = PyanNet.from_pretrained('20251208_Sonogram_Segmentation.ckpt')
+segModel.specifications = newSpecs
+segmentation_duration = segModel.specifications.duration
+
+# Update baseline with training
+pipeline._segmentation = Inference(
+    segModel,
+    duration=segmentation_duration,
+    step=pipeline.segmentation_step * segmentation_duration,
+    skip_aggregation=True,
+    batch_size=1,
+)
+
 pipeline.to(device)
 
 # ---------------------------------------------------------------------------