"""
utils.py — pure logic helpers with no Streamlit dependency.
Covers:
- Audio processing (processFile, clip extraction, randomization)
- DataFrame builders (build_df2 … build_df5)
- Plotly figure builders (one function per chart tab)
- Multi-file summary DataFrame builders
"""
import io
import random
import datetime as dt
import copy
import numpy as np
import pandas as pd
import soundfile as sf
import torch
import plotly.express as px
import plotly.graph_objects as go
import sonogram_utility as su
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
CLIP_MIN_S = 3.0
CLIP_MAX_S = 5.0
CLIP_SCAN_STEP_S = 0.5
TRANSPARENT_BG = dict(
plot_bgcolor="rgba(0,0,0,0)",
paper_bgcolor="rgba(0,0,0,0)",
)
# ---------------------------------------------------------------------------
# Audio processing
# ---------------------------------------------------------------------------
def processFile(filePath, pipeline, enableDenoise, earlyCleanup,
gainWindow, minimumGain, maximumGain,
dfModel=None, dfState=None, attenLimDB=3):
"""Load, optionally denoise, equalize, and diarize an audio file.
Returns (annotations, totalTimeInSeconds, waveform_tensor, sample_rate).
"""
print("Loading file")
waveformList, sampleRate = su.splitIntoTimeSegments(filePath, 600)
print("File loaded")
enhancedWaveformList = []
if enableDenoise:
print("Denoising")
for w in waveformList:
if enableDenoise:
from df import enhance
newW = enhance(dfModel, dfState, w, atten_lim_db=attenLimDB).detach().cpu()
enhancedWaveformList.append(newW)
else:
enhancedWaveformList.append(w)
if enableDenoise:
print("Audio denoised")
waveformEnhanced = su.combineWaveforms(enhancedWaveformList)
if earlyCleanup:
del enhancedWaveformList
print("Equalizing Audio")
waveform_gain_adjusted = su.equalizeVolume()(
waveformEnhanced, sampleRate, gainWindow, minimumGain, maximumGain
)
if earlyCleanup:
del waveformEnhanced
print("Audio Equalized")
print("Detecting speakers")
diarization_output = pipeline({"waveform": waveform_gain_adjusted, "sample_rate": sampleRate})
annotations = diarization_output.speaker_diarization
print("Speakers Detected")
totalTimeInSeconds = int(waveform_gain_adjusted.shape[-1] / sampleRate)
return annotations, totalTimeInSeconds, waveform_gain_adjusted, sampleRate
# ---------------------------------------------------------------------------
# Speaker sample helpers
# ---------------------------------------------------------------------------
def extract_clip_bytes(waveform, sample_rate, seg_start, seg_end):
"""Return WAV bytes for the loudest SAMPLE_MIN–SAMPLE_MAX window in [seg_start, seg_end]."""
total_samples = waveform.shape[-1]
seg_start_s = int(seg_start * sample_rate)
seg_end_s = min(int(seg_end * sample_rate), total_samples)
seg_dur = (seg_end_s - seg_start_s) / sample_rate
clip_dur = min(max(min(seg_dur, CLIP_MAX_S), CLIP_MIN_S), seg_dur)
clip_samples = int(clip_dur * sample_rate)
best_start = seg_start_s
best_rms = -1.0
step_samples = int(CLIP_SCAN_STEP_S * sample_rate)
pos = seg_start_s
while pos + clip_samples <= seg_end_s:
window = waveform[:, pos: pos + clip_samples].float()
rms = float(window.pow(2).mean().sqrt())
if rms > best_rms:
best_rms = rms
best_start = pos
pos += step_samples
clip_np = waveform[:, best_start: best_start + clip_samples].numpy().T
buf = io.BytesIO()
sf.write(buf, clip_np, sample_rate, format="WAV", subtype="PCM_16")
buf.seek(0)
return buf.read()
def build_speaker_clips(annotations, waveform, sample_rate):
"""Return (samples_dict, segments_dict) for all speakers in annotations.
samples_dict : {speaker: wav_bytes}
segments_dict : {speaker: [(start, end), ...]}
"""
clips = {}
segments = {}
for speaker in annotations.labels():
speaker_segments = [
seg for seg, _, label in annotations.itertracks(yield_label=True)
if label == speaker
]
if not speaker_segments:
continue
segments[speaker] = [(s.start, s.end) for s in speaker_segments]
longest = max(speaker_segments, key=lambda s: s.duration)
clips[speaker] = extract_clip_bytes(waveform, sample_rate, longest.start, longest.end)
return clips, segments
def get_randomized_clip(waveform, sample_rate, segments):
"""Return WAV bytes for a random 3–5 s audio sample drawn from a random segment.
segments : [(start, end), ...] (all segments for one speaker)
"""
durations = [max(e - s, 0.01) for s, e in segments]
total_dur = sum(durations)
rand_val = random.random() * total_dur
cumulative = 0.0
chosen_start, chosen_end = segments[0]
for (seg_s, seg_e), dur in zip(segments, durations):
cumulative += dur
if rand_val <= cumulative:
chosen_start, chosen_end = seg_s, seg_e
break
seg_dur = chosen_end - chosen_start
clip_dur = min(max(min(seg_dur, CLIP_MAX_S), CLIP_MIN_S), seg_dur)
max_offset = max(seg_dur - clip_dur, 0.0)
offset = random.uniform(0.0, max_offset)
clip_start = chosen_start + offset
clip_end = clip_start + clip_dur
return extract_clip_bytes(waveform, sample_rate, clip_start, clip_end)
# ---------------------------------------------------------------------------
# DataFrame builders (called from analyze() in state.py)
# ---------------------------------------------------------------------------
def build_df3(noVoice, oneVoice, multiVoice):
"""Voice category totals DataFrame."""
return pd.DataFrame({
"values": [su.sumTimes(noVoice), su.sumTimes(oneVoice), su.sumTimes(multiVoice)],
"names": ["No Voice", "Single Voice", "Multi Voice"],
})
def build_df4(speakerNames, categorySelections, categoryNames, currAnnotation):
"""Speaker-to-category time DataFrame. Returns (df4, nameList, valueList, extraNames, extraValues)."""
nameList = list(categoryNames)
valueList = [0.0] * len(nameList)
extraNames : list = []
extraValues: list = []
for sp in speakerNames:
found = False
for i, _ in enumerate(nameList):
if sp in categorySelections[i]:
valueList[i] += su.sumTimes(currAnnotation.subset([sp]))
found = True
break
if not found:
extraNames.append(sp)
extraValues.append(su.sumTimes(currAnnotation.subset([sp])))
if extraNames:
pairs = sorted(zip(extraNames, extraValues), key=lambda p: p[0])
extraNames, extraValues = map(list, zip(*pairs))
else:
extraNames, extraValues = [], []
df4 = pd.DataFrame({"values": valueList + extraValues, "names": nameList + extraNames})
return df4, nameList, valueList, extraNames, extraValues
def build_df5(oneVoice, multiVoice, sumNoVoice, sumOneVoice, sumMultiVoice, currTotalTime):
"""Hierarchical voice-category DataFrame for sunburst / treemap."""
speakerList, timeList = su.sumTimesPerSpeaker(oneVoice)
multiSpeakerList, multiTimeList = su.sumMultiTimesPerSpeaker(multiVoice)
speakerList = list(speakerList) if speakerList else []
timeList = list(timeList) if timeList else []
multiSpeakerList = list(multiSpeakerList) if multiSpeakerList else []
multiTimeList = list(multiTimeList) if multiTimeList else []
summativeMulti = sum(multiTimeList) if multiTimeList else 1
safeOneVoice = sumOneVoice if sumOneVoice > 0 else 1
safeTotalTime = currTotalTime if currTotalTime > 0 else 1
base = [sumNoVoice / safeTotalTime, sumOneVoice / safeTotalTime, sumMultiVoice / safeTotalTime]
timeStrings = su.timeToString(timeList) if timeList else []
multiTimeStrings = su.timeToString(multiTimeList) if multiTimeList else []
if isinstance(timeStrings, str):
timeStrings = [timeStrings]
if isinstance(multiTimeStrings, str):
multiTimeStrings = [multiTimeStrings]
n_ov = len(speakerList)
n_mv = len(multiSpeakerList)
return pd.DataFrame({
"ids": ["NV", "OV", "MV"] + [f"OV_{i}" for i in range(n_ov)] + [f"MV_{i}" for i in range(n_mv)],
"labels": ["No Voice", "Single Voice", "Multi Voice"] + speakerList + multiSpeakerList,
"parents": ["", "", ""] + ["OV"] * n_ov + ["MV"] * n_mv,
"parentNames": ["Total", "Total", "Total"] + ["Single Voice"] * n_ov + ["Multi Voice"] * n_mv,
"values": [sumNoVoice, sumOneVoice, sumMultiVoice] + timeList + multiTimeList,
"valueStrings": [
su.timeToString(sumNoVoice),
su.timeToString(sumOneVoice),
su.timeToString(sumMultiVoice),
] + timeStrings + multiTimeStrings,
"percentiles": [b * 100 for b in base]
+ [(t * 100) / safeOneVoice * base[1] for t in timeList]
+ [(t * 100) / summativeMulti * base[2] for t in multiTimeList],
"parentPercentiles": [b * 100 for b in base]
+ [(t * 100) / safeOneVoice for t in timeList]
+ [(t * 100) / summativeMulti for t in multiTimeList],
})
def build_df2(df4_names, df4_values, currTotalTime):
"""Time-spoken DataFrame (raw seconds) used by the bar chart tab."""
return pd.DataFrame({
"values": list(df4_values),
"names": df4_names,
})
# ---------------------------------------------------------------------------
# Plotly figure builders
# ---------------------------------------------------------------------------
def _save_fig(fig, *paths):
"""Try to write fig to each path; silently skip on failure."""
for path in paths:
try:
fig.write_image(path)
except Exception:
pass
def build_fig_pie1(df3, catTypeColors):
"""Voice category pie chart."""
fig = go.Figure()
fig.update_layout(
title_text="Percentage of each Voice Category",
colorway=catTypeColors,
**TRANSPARENT_BG,
)
fig.add_trace(go.Pie(values=df3["values"], labels=df3["names"], sort=False))
return fig
def build_fig_pie2(df4, speakerNames, speakerColors, catColors, get_display_name_fn, currFile):
"""Speaker / category pie chart."""
df4 = df4.copy()
figColors = [
speakerColors[list(speakerNames).index(n)]
for n in df4["names"] if n in speakerNames
]
df4["names"] = df4["names"].apply(lambda s: get_display_name_fn(s, currFile))
fig = go.Figure()
fig.update_layout(
title_text="Percentage of Speakers per Role",
colorway=catColors + figColors,
**TRANSPARENT_BG,
)
fig.add_trace(go.Pie(values=df4["values"], labels=df4["names"], sort=False))
return fig
def build_fig_sunburst(df5, catTypeColors, speakerColors, get_display_name_fn, currFile):
"""Sunburst voice-category chart."""
df5 = df5.copy()
df5["labels"] = df5["labels"].apply(lambda s: get_display_name_fn(s, currFile))
df5["parentNames"] = df5["parentNames"].apply(lambda s: get_display_name_fn(s, currFile))
fig = px.sunburst(
df5,
branchvalues="total",
names="labels", ids="ids", parents="parents",
values="percentiles",
custom_data=["labels", "valueStrings", "percentiles", "parentNames", "parentPercentiles"],
color="labels",
title="Percentage of each Voice Category with Speakers",
color_discrete_sequence=catTypeColors + speakerColors,
)
fig.update_traces(hovertemplate="
".join([
"%{customdata[0]}",
"Duration: %{customdata[1]}s",
"Percentage of Total: %{customdata[2]:.2f}%",
"Parent: %{customdata[3]}",
"Percentage of Parent: %{customdata[4]:.2f}%",
]))
fig.update_layout(**TRANSPARENT_BG)
return fig
def build_fig_treemap(df5, catTypeColors, speakerColors, get_display_name_fn, currFile):
"""Treemap voice-category chart."""
df5 = df5.copy()
df5["labels"] = df5["labels"].apply(lambda s: get_display_name_fn(s, currFile))
df5["parentNames"] = df5["parentNames"].apply(lambda s: get_display_name_fn(s, currFile))
fig = px.treemap(
df5,
branchvalues="total",
names="labels", parents="parents", ids="ids",
values="percentiles",
custom_data=["labels", "valueStrings", "percentiles", "parentNames", "parentPercentiles"],
color="labels",
title="Division of Speakers in each Voice Category",
color_discrete_sequence=catTypeColors + speakerColors,
)
fig.update_traces(hovertemplate="
".join([
"%{customdata[0]}",
"Duration: %{customdata[1]}s",
"Percentage of Total: %{customdata[2]:.2f}%",
"Parent: %{customdata[3]}",
"Percentage of Parent: %{customdata[4]:.2f}%",
]))
fig.update_layout(**TRANSPARENT_BG)
return fig
def build_fig_timeline(speakers_dataFrame, currTotalTime, speakerColors, get_display_name_fn, currFile):
"""Gantt-style speaker timeline."""
df = speakers_dataFrame.copy()
df["Resource"] = df["Resource"].apply(lambda s: get_display_name_fn(s, currFile))
base = dt.datetime.combine(dt.date.today(), dt.time.min)
def to_audio_dt(s):
if isinstance(s, (dt.datetime, pd.Timestamp)):
midnight = s.replace(hour=0, minute=0, second=0, microsecond=0)
seconds = (s - midnight).total_seconds()
else:
seconds = float(s)
return base + dt.timedelta(seconds=seconds)
df["Start"] = df["Start"].apply(to_audio_dt)
df["Finish"] = df["Finish"].apply(to_audio_dt)
fig = px.timeline(
df, x_start="Start", x_end="Finish", y="Resource", color="Resource",
title="Timeline of Audio with Speakers",
color_discrete_sequence=speakerColors,
)
fig.update_yaxes(autorange="reversed")
h = int(currTotalTime // 3600)
m = int(currTotalTime % 3600 // 60)
s = int(currTotalTime % 60)
ms= int(currTotalTime * 1_000_000 % 1_000_000)
time_max = dt.time(h, m, s, ms)
fig.update_layout(
xaxis_tickformatstops=[
dict(dtickrange=[None, 1000], value="%H:%M:%S.%L"),
dict(dtickrange=[1000, None], value="%H:%M:%S"),
],
xaxis=dict(range=[
dt.datetime.combine(dt.date.today(), dt.time.min),
dt.datetime.combine(dt.date.today(), time_max),
]),
xaxis_title="Time",
yaxis_title=None,
showlegend=False,
yaxis={"showticklabels": True},
**TRANSPARENT_BG,
)
return fig
def _seconds_to_hhmmss(seconds):
"""Convert a float seconds value to a hh:mm:ss.ss string."""
seconds = float(seconds)
h = int(seconds // 3600)
m = int((seconds % 3600) // 60)
s = seconds % 60
return f"{h:02d}:{m:02d}:{s:05.2f}"
def build_fig_bar(df2, speakerNames, catColors, speakerColors, get_display_name_fn, currFile):
"""Horizontal bar chart — time spoken per speaker (hh:mm:ss.ss).
Only individual speakers are shown; role/category rows are excluded.
"""
df2 = df2.copy()
# Keep only rows whose raw name is an actual speaker (not a category label)
df2 = df2[df2["names"].isin(speakerNames)]
df2["names"] = df2["names"].apply(lambda s: get_display_name_fn(s, currFile))
df2["time_label"] = df2["values"].apply(_seconds_to_hhmmss)
fig = px.bar(
df2, x="values", y="names", color="names", orientation="h",
custom_data=["names", "time_label"],
title="Time Spoken by each Speaker",
color_discrete_sequence=catColors + speakerColors,
)
# Hide x-axis tick labels — values are crowded with many speakers.
# The exact time is still visible on hover via the hovertemplate.
fig.update_xaxes(showticklabels=False)
fig.update_yaxes(autorange="reversed")
fig.update_layout(
xaxis_title="Time Spoken",
yaxis_title=None,
showlegend=False,
yaxis={"showticklabels": True},
**TRANSPARENT_BG,
)
fig.update_traces(hovertemplate="
".join([
"%{customdata[0]}",
"Time Spoken: %{customdata[1]}",
]))
return fig
# ---------------------------------------------------------------------------
# Multi-file summary DataFrames
# ---------------------------------------------------------------------------
def build_multifile_category_df(validNames, results, summaries, categories, categorySelect,
speakerRenames=None):
"""Build df6 (category breakdown per file) for the multi-file expander.
Uses su.sumTimes() per speaker (same as the single-file charts) so that:
- Each speaker's time = union of their segments (overlaps within one speaker
are merged by get_timeline().duration())
- Multiple speakers in the same role are subset-unioned before summing so
cross-speaker overlaps within a role are counted only once
- Values are proportions (0-1) of the file's total duration
speakerRenames: {filename: {raw_sp: display_name}} — applied to unassigned
speaker column headers.
"""
speakerRenames = speakerRenames or {}
df6_dict = {"files": validNames}
allCategories = copy.deepcopy(categories)
# First pass: discover unassigned speaker columns across all files
for fn in validNames:
currAnnotation, _ = results[fn]
prefix = fn + ": "
assigned = {
t[len(prefix):]
for tokens in categorySelect
for t in tokens
if t.startswith(prefix)
}
renames = speakerRenames.get(fn, {})
for sp in currAnnotation.labels():
if sp not in assigned:
display = renames.get(sp, sp)
if display not in allCategories:
allCategories.append(display)
df6_dict.setdefault(display, [])
for category in categories:
df6_dict.setdefault(category, [])
# Second pass: compute proportions per file
for row_idx, fn in enumerate(validNames):
currAnnotation, totalSeconds = results[fn]
safe_total = max(totalSeconds, 1)
prefix = fn + ": "
renames = speakerRenames.get(fn, {})
# Track which allCategories columns get a value this row
filled = set()
# For each role: union all assigned speakers into one subset, then sum.
# su.sumTimes uses get_timeline(False).duration() which merges overlaps.
for i, category in enumerate(categories):
assigned_sps = [
t[len(prefix):]
for t in categorySelect[i]
if t.startswith(prefix)
] if i < len(categorySelect) else []
valid_sps = [sp for sp in assigned_sps if sp in currAnnotation.labels()]
if valid_sps:
val = su.sumTimes(currAnnotation.subset(valid_sps)) / safe_total
else:
val = 0.0
df6_dict[category].append(min(val, 1.0))
filled.add(category)
# For unassigned speakers: each gets their own column
assigned_all = {
t[len(prefix):]
for tokens in categorySelect
for t in tokens
if t.startswith(prefix)
}
unassigned = [sp for sp in currAnnotation.labels() if sp not in assigned_all]
for sp in unassigned:
display = renames.get(sp, sp)
val = su.sumTimes(currAnnotation.subset([sp])) / safe_total
df6_dict[display].append(min(val, 1.0))
filled.add(display)
# Fill 0 for every allCategories column not touched this row
for category in allCategories:
if category not in filled:
df6_dict[category].append(0)
return pd.DataFrame(df6_dict), allCategories
def build_multifile_role_voice_df(validNames, results, summaries, categories,
categorySelect, speakerRenames=None):
"""Build df8: per-file proportions split by role for single voice, plus
Multi Voice and No Voice.
Single Voice time is broken down into each role and an Unassigned bucket
(speakers in single-voice segments that haven't been assigned to any role).
Multi Voice and No Voice come from df5 percentiles (0-100 scale) converted
to 0-1 proportions.
This is the combination of df6 (role proportions) and df7 (voice categories)
where Single Voice is replaced by its constituent roles.
"""
speakerRenames = speakerRenames or {}
col_names = list(categories) + ["Unassigned", "Multi Voice", "No Voice"]
df8_dict = {"files": validNames}
for col in col_names:
df8_dict[col] = []
for fn in validNames:
currAnnotation, totalSeconds = results[fn]
safe_total = max(totalSeconds, 1)
prefix = fn + ": "
renames = speakerRenames.get(fn, {})
# Role proportions — same logic as build_multifile_category_df
assigned_all = set()
for i, category in enumerate(categories):
assigned_sps = [
t[len(prefix):]
for t in (categorySelect[i] if i < len(categorySelect) else [])
if t.startswith(prefix)
]
valid_sps = [sp for sp in assigned_sps if sp in currAnnotation.labels()]
assigned_all.update(valid_sps)
if valid_sps:
val = su.sumTimes(currAnnotation.subset(valid_sps)) / safe_total
else:
val = 0.0
df8_dict[category].append(min(val, 1.0))
# Unassigned speakers
unassigned_sps = [sp for sp in currAnnotation.labels() if sp not in assigned_all]
if unassigned_sps:
val = su.sumTimes(currAnnotation.subset(unassigned_sps)) / safe_total
else:
val = 0.0
df8_dict["Unassigned"].append(min(val, 1.0))
# Multi Voice and No Voice from df5 percentiles (0-100 → 0-1)
partial = summaries[fn]["df5"]
df8_dict["No Voice"].append(partial["percentiles"][0] / 100)
df8_dict["Multi Voice"].append(partial["percentiles"][2] / 100)
return pd.DataFrame(df8_dict), col_names
def build_multifile_voice_df(validNames, summaries):
"""Build df7 (no/one/multi voice percentages per file) for the multi-file expander."""
voiceNames = ["No Voice", "Single Voice", "Multi Voice"]
df7_dict = {"files": validNames}
for name in voiceNames:
df7_dict[name] = []
for fn in validNames:
partial = summaries[fn]["df5"]
for i, name in enumerate(voiceNames):
df7_dict[name].append(partial["percentiles"][i])
return pd.DataFrame(df7_dict), voiceNames