#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
segmentation_infer_smooth_segments.py
- Loads WhisperOddEven checkpoint
/home/user/outs/segmentation_gemini_2p_medium_model_best.pt
(override via CKPT env var).
- For each audio file in AUDIO_INPUT_DIR:
* load, resample to 16 kHz mono
* split into 30 s chunks
* run segmentation
* SMOOTH each track so that no segment (incl. background 0) is shorter than
MIN_SEGMENT_SEC seconds
* extract per-track segments (odd/even) and cut audio snippets
* build a MERGED timeline that starts/ends segments whenever either track
changes label, then smooth that merged timeline so that each merged
segment is also at least MIN_SEGMENT_SEC long, merging short segments
with neighbors using the rules described below.
- Writes a single HTML report with:
* smoothed per-track heatmap
* merged-timeline heatmap
* tables of per-track segments (with audio players)
* tables of merged segments (with audio players)
Merging rule for short merged segments:
- If a merged segment is shorter than MIN_SEGMENT_SEC, merge it with one of its
immediate neighbors.
- Prefer the neighbor whose (odd_label, even_label) matches this segment best
(majority vote over the two labels).
- If similarity is equal (or one neighbor is missing), merge with the neighbor
that has the shorter duration. If still equal, merge with the left neighbor.
"""
from __future__ import annotations
import os
import io
import sys
import time
import math
import base64
import shutil
from pathlib import Path
from typing import List, Dict, Any, Tuple
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# plotting
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
# audio
import soundfile as sf
import librosa
from pydub import AudioSegment # requires ffmpeg
from transformers import WhisperFeatureExtractor, WhisperModel
# =========================
# ========== CONFIG =======
# =========================
AUDIO_INPUT_DIR = Path(os.getenv("AUDIO_INPUT_DIR", "./infer-audio"))
OUT_DIR = Path(os.getenv("OUT_DIR", "./outs_infer"))
CKPT_PATH = Path(os.getenv("CKPT", "/home/user/outs/segmentation_gemini_medium_no_overlap_4epochs_model_best.pt"))
HF_MODEL_ID = os.getenv("HF_MODEL_ID", "openai/whisper-small")
USE_LOCAL_MODELS = bool(int(os.getenv("USE_LOCAL_MODELS", "0")))
MODELS_SNAPSHOT_DIR = Path(os.getenv("MODELS_SNAPSHOT_DIR", "")) if USE_LOCAL_MODELS else None
HF_HOME = Path(os.getenv("HF_HOME", (OUT_DIR / ".hf")))
TRANSFORMERS_CACHE = Path(os.getenv("TRANSFORMERS_CACHE", (OUT_DIR / ".hf" / "hub")))
MIXED_PRECISION = os.getenv("MIXED_PRECISION", "auto").lower()
# constants (must match training)
SAMPLE_RATE = 16000
CLIP_SECONDS = 30.0
NUM_FRAMES = 1500
NUM_TRACKS = 2
MAX_SEGMENTS = 20
# --- MINIMUM SEGMENT LENGTH (seconds) for both per-track and merged segments ---
MIN_SEGMENT_SEC = float(os.getenv("MIN_SEGMENT_SEC", "1.0"))
MIN_SEGMENT_FRAMES = max(1, int(round(MIN_SEGMENT_SEC * NUM_FRAMES / CLIP_SECONDS)))
FFMPEG_AVAILABLE = shutil.which("ffmpeg") is not None
WARNED_NO_FFMPEG = False
# =========================
# ====== BASIC SETUP ======
# =========================
def setup_dirs():
OUT_DIR.mkdir(parents=True, exist_ok=True)
(OUT_DIR / ".mplconfig").mkdir(parents=True, exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str((OUT_DIR / ".mplconfig").resolve()))
HF_HOME.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HF_HOME", str(HF_HOME.resolve()))
os.environ.setdefault("TRANSFORMERS_CACHE", str(TRANSFORMERS_CACHE.resolve()))
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,max_split_size_mb:128")
def preferred_dtype():
if MIXED_PRECISION == "bf16":
return torch.bfloat16
if MIXED_PRECISION == "fp16":
return torch.float16
if MIXED_PRECISION == "fp32":
return torch.float32
if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
return torch.bfloat16
return torch.float16 if torch.cuda.is_available() else torch.float32
def _model_resolved_name(model_id: str) -> Tuple[str, bool]:
if USE_LOCAL_MODELS and MODELS_SNAPSHOT_DIR and MODELS_SNAPSHOT_DIR.is_dir():
local_dirname = model_id.replace("/", "__")
cand = MODELS_SNAPSHOT_DIR / local_dirname
if cand.is_dir():
return str(cand), True
return model_id, False
# =========================
# ========= MODEL =========
# =========================
class WhisperOddEven(nn.Module):
def __init__(self, base_id: str, freeze_encoder: bool = False):
super().__init__()
resolved, is_local = _model_resolved_name(base_id)
self.whisper = WhisperModel.from_pretrained(resolved, local_files_only=is_local)
# decoder unused
for p in self.whisper.decoder.parameters():
p.requires_grad = False
for p in self.whisper.encoder.parameters():
p.requires_grad = not freeze_encoder
d_model = self.whisper.config.d_model
hidden = max(256, d_model // 2)
self.head = nn.Sequential(
nn.Linear(d_model, hidden),
nn.GELU(),
nn.Linear(hidden, NUM_TRACKS * (MAX_SEGMENTS + 1)),
)
def forward(self, input_features: torch.FloatTensor):
enc = self.whisper.encoder(input_features=input_features).last_hidden_state # [B,1500,D]
B, T, D = enc.shape
logits = self.head(enc) # [B,T,NUM_TRACKS*(C)]
C = MAX_SEGMENTS + 1
logits = logits.view(B, T, NUM_TRACKS, C).permute(0, 2, 1, 3).contiguous()
return logits # [B,2,1500,C]
# =========================
# ====== AUDIO UTILS ======
# =========================
def load_audio_mono_16k(path: Path) -> np.ndarray:
wav, sr = librosa.load(str(path), sr=SAMPLE_RATE, mono=True)
if wav.ndim > 1:
wav = wav.mean(axis=0)
return wav.astype(np.float32, copy=False)
def split_into_chunks(wav: np.ndarray, sr: int, clip_seconds: float):
chunk_size = int(clip_seconds * sr)
total = len(wav)
if total == 0:
return []
n_chunks = math.ceil(total / chunk_size)
chunks = []
for i in range(n_chunks):
start = i * chunk_size
end = min(start + chunk_size, total)
seg = wav[start:end]
if len(seg) < chunk_size:
seg = np.pad(seg, (0, chunk_size - len(seg)), mode="constant")
chunks.append((i, start, seg.astype(np.float32, copy=False)))
return chunks
def wav_chunk_to_audio_bytes(wav: np.ndarray, sr: int):
"""
Try to export as MP3 (if ffmpeg is available). Otherwise fall back to WAV.
Returns (audio_bytes, mime_type).
"""
global WARNED_NO_FFMPEG
buf_wav = io.BytesIO()
sf.write(buf_wav, wav, sr, format="WAV")
wav_bytes = buf_wav.getvalue()
if not FFMPEG_AVAILABLE:
if not WARNED_NO_FFMPEG:
print("[audio] ffmpeg not found; embedding WAV instead of MP3.", flush=True)
WARNED_NO_FFMPEG = True
return wav_bytes, "audio/wav"
try:
buf_wav.seek(0)
audio = AudioSegment.from_file(buf_wav, format="wav")
out_buf = io.BytesIO()
audio.export(out_buf, format="mp3", bitrate="128k")
out_buf.seek(0)
return out_buf.read(), "audio/mpeg"
except Exception as e:
if not WARNED_NO_FFMPEG:
print(f"[audio] Failed to encode MP3, falling back to WAV: {e}", flush=True)
WARNED_NO_FFMPEG = True
return wav_bytes, "audio/wav"
# =========================
# ====== SEGMENT OPS ======
# =========================
def smooth_min_duration(ids: np.ndarray, min_frames: int, max_iter: int = 10) -> np.ndarray:
"""
Enforce a minimum run length (in frames) for an ID sequence (1D).
Shorter runs are reassigned to the longer of their neighbors, iteratively.
"""
ids = ids.copy()
n = len(ids)
if n == 0:
return ids
for _ in range(max_iter):
runs = []
start = 0
cur = ids[0]
for i in range(1, n):
if ids[i] != cur:
runs.append((cur, start, i))
start = i
cur = ids[i]
runs.append((cur, start, n))
changed = False
for ri, (label, s, e) in enumerate(runs):
length = e - s
if length >= min_frames:
continue
left = runs[ri - 1] if ri > 0 else None
right = runs[ri + 1] if ri + 1 < len(runs) else None
if left is None and right is None:
continue
if left is None:
new_label = right[0]
elif right is None:
new_label = left[0]
else:
len_left = left[2] - left[1]
len_right = right[2] - right[1]
new_label = left[0] if len_left >= len_right else right[0]
if new_label != label:
ids[s:e] = new_label
changed = True
if not changed:
break
return ids
def extract_segments(ids: np.ndarray, include_bg: bool = False):
"""
Return list of (label, frame_start, frame_end) runs.
Optionally filter out background label 0.
"""
n = len(ids)
if n == 0:
return []
runs = []
start = 0
cur = ids[0]
for i in range(1, n):
if ids[i] != cur:
runs.append((cur, start, i))
start = i
cur = ids[i]
runs.append((cur, start, n))
if not include_bg:
runs = [(lab, s, e) for (lab, s, e) in runs if lab != 0]
return runs
def frames_to_times(s: int, e: int):
start_t = s / NUM_FRAMES * CLIP_SECONDS
end_t = e / NUM_FRAMES * CLIP_SECONDS
return start_t, end_t
def cut_wav(seg_wav: np.ndarray, start_t: float, end_t: float) -> np.ndarray:
start_samp = int(round(start_t * SAMPLE_RATE))
end_samp = int(round(end_t * SAMPLE_RATE))
start_samp = max(0, min(start_samp, len(seg_wav)))
end_samp = max(start_samp + 1, min(end_samp, len(seg_wav)))
return seg_wav[start_samp:end_samp]
# =========================
# ==== MERGED TIMELINE ====
# =========================
def smooth_merged_segments(merged: List[Tuple[int,int,int,int]], min_frames: int) -> List[Tuple[int,int,int,int]]:
"""
Enforce minimum length for merged segments.
merged: list of (frame_start, frame_end, odd_label, even_label).
If a segment has length < min_frames, we merge it with a neighbor:
- If both neighbors exist, choose the one with higher similarity of
(odd_label, even_label). Similarity is number of matching labels (0..2).
- If similarity is equal, merge with the neighbor that has shorter
duration (in frames). If still equal, merge with the left neighbor.
- If only one neighbor exists, merge with that neighbor.
Returns a new merged list.
"""
if len(merged) <= 1:
return merged
merged = list(merged)
def seg_len(seg):
return seg[1] - seg[0]
def sim(a, b):
# a,b: (fs,fe, odd,even)
score = 0
if a[2] == b[2]:
score += 1
if a[3] == b[3]:
score += 1
return score
changed = True
while changed:
changed = False
n = len(merged)
if n <= 1:
break
for i, seg in enumerate(merged):
length = seg_len(seg)
if length >= min_frames:
continue
left = merged[i - 1] if i > 0 else None
right = merged[i + 1] if i + 1 < n else None
if left is None and right is None:
continue
# Decide which neighbor to merge with
if left is not None and right is not None:
s_left = sim(seg, left)
s_right = sim(seg, right)
if s_left > s_right:
target = "left"
elif s_right > s_left:
target = "right"
else:
# similarity tie -> choose shorter neighbor
len_left = seg_len(left)
len_right = seg_len(right)
if len_left < len_right:
target = "left"
elif len_right < len_left:
target = "right"
else:
target = "left" # full tie -> left
elif left is not None:
target = "left"
else:
target = "right"
if target == "left":
fs = left[0]
fe = seg[1]
odd_label = left[2]
even_label = left[3]
merged[i - 1] = (fs, fe, odd_label, even_label)
del merged[i]
else:
fs = seg[0]
fe = right[1]
odd_label = right[2]
even_label = right[3]
merged[i + 1] = (fs, fe, odd_label, even_label)
del merged[i]
changed = True
break # restart scanning with new list
return merged
def build_merged_segments(ids_odd: np.ndarray, ids_even: np.ndarray, min_frames: int):
"""
Build merged segmentation from two tracks and then smooth merged segments.
- boundaries are at 0, NUM_FRAMES, and every point where either track changes.
- for each raw merged segment we set odd/even labels via majority label.
- then we enforce minimum length for the merged segments via
smooth_merged_segments.
"""
assert len(ids_odd) == len(ids_even) == NUM_FRAMES
n = NUM_FRAMES
boundaries = {0, n}
for ids in (ids_odd, ids_even):
cur = ids[0]
for i in range(1, n):
if ids[i] != cur:
boundaries.add(i)
cur = ids[i]
b = sorted(boundaries)
merged = []
for i in range(len(b) - 1):
s = b[i]
e = b[i + 1]
if e <= s:
continue
slice_odd = ids_odd[s:e]
slice_even = ids_even[s:e]
if slice_odd.size == 0 or slice_even.size == 0:
continue
odd_vals, odd_counts = np.unique(slice_odd, return_counts=True)
even_vals, even_counts = np.unique(slice_even, return_counts=True)
odd_label = int(odd_vals[np.argmax(odd_counts)])
even_label = int(even_vals[np.argmax(even_counts)])
merged.append((s, e, odd_label, even_label))
# Now enforce min length also on merged segments
merged = smooth_merged_segments(merged, min_frames)
return merged
# =========================
# ======= PLOTTING ========
# =========================
def _plot_tracks_seconds(pred_ids: torch.Tensor, title: str) -> bytes:
"""
pred_ids: [2, NUM_FRAMES] LongTensor
"""
secs = np.linspace(0.0, CLIP_SECONDS, NUM_FRAMES)
fig = plt.figure(figsize=(10, 2.8))
ax = plt.gca()
im = ax.imshow(
pred_ids.numpy(),
aspect="auto",
interpolation="nearest",
origin="upper",
extent=[secs[0], secs[-1], -0.5, 1.5],
)
ax.set_title(title)
ax.set_xlabel("Time (s)")
ax.set_yticks([0, 1])
ax.set_yticklabels(["odd", "even"])
cb = plt.colorbar(im, fraction=0.046, pad=0.04)
cb.set_label("Segment ID")
buf = io.BytesIO()
fig.savefig(buf, format="png", dpi=150, bbox_inches="tight")
plt.close(fig)
buf.seek(0)
return buf.read()
def _plot_merged_segments(seg_ids: np.ndarray, title: str) -> bytes:
"""
seg_ids: [NUM_FRAMES] array where each frame holds a merged-segment index.
"""
secs = np.linspace(0.0, CLIP_SECONDS, NUM_FRAMES)
fig = plt.figure(figsize=(10, 2.8))
ax = plt.gca()
im = ax.imshow(
seg_ids[np.newaxis, :],
aspect="auto",
interpolation="nearest",
origin="upper",
extent=[secs[0], secs[-1], -0.5, 0.5],
)
ax.set_title(title)
ax.set_xlabel("Time (s)")
ax.set_yticks([0])
ax.set_yticklabels(["merged"])
cb = plt.colorbar(im, fraction=0.046, pad=0.04)
cb.set_label("Merged seg ID")
buf = io.BytesIO()
fig.savefig(buf, format="png", dpi=150, bbox_inches="tight")
plt.close(fig)
buf.seek(0)
return buf.read()
# =========================
# ========= HTML ==========
# =========================
def write_html_report(out_dir: Path, chunks: List[Dict[str, Any]]) -> Path:
ts = time.strftime("%Y%m%d_%H%M%S")
html = [f"""
Odd/Even Segmentation - Inference {ts}
Odd/Even Segmentation - Inference
This report shows smoothed segmentations for each 30-second chunk of your audio files.
The model predicts two parallel time tracks ("odd" and "even") that can hold overlapping events.
We first smooth each track so that no segment (including background 0) is shorter than {MIN_SEGMENT_SEC:.2f} seconds.
Then:
Per-track segments: segments for each track (odd/even) with duration >= {MIN_SEGMENT_SEC:.2f}s, each with its own audio player.
Merged timeline: a single segmentation where a new segment starts or ends whenever either track changes, and each merged segment is also at least {MIN_SEGMENT_SEC:.2f}s long by merging very short segments into their most similar neighbor.
"""]
for ch in chunks:
html.append(f"""
{ch['file_name']} - chunk {ch['chunk_idx']}
Chunk offset in file: {ch['chunk_offset']:.2f} - {ch['chunk_offset'] + CLIP_SECONDS:.2f} s
Smoothed per-track predictions (odd/even).Merged timeline: segment borders whenever odd or even track changes label, then smoothed to enforce a minimum duration.
Per-track segments (min {MIN_SEGMENT_SEC:.2f} s)
Each row is one predicted event on the odd or even track. Times are relative to the start of this 30-second chunk.
#
Track
Label ID
Start (s)
End (s)
Duration (s)
Audio
""")
# per-track table
for i, seg in enumerate(ch["track_segments"], start=1):
audio_cell = ""
if seg["audio_b64"] and seg["audio_mime"]:
audio_cell = (
''
)
html.append(
f"
{i}
"
f"
{seg['track']}
"
f"
{seg['label']}
"
f"
{seg['start']:.2f}
"
f"
{seg['end']:.2f}
"
f"
{seg['dur']:.2f}
"
f"
{audio_cell}
"
)
html.append("
")
# merged timeline table
html.append(f"""
Merged timeline segments
The merged timeline splits the 30-second chunk wherever either the odd or even track changes label.
Very short merged segments (shorter than {MIN_SEGMENT_SEC:.2f}s) are merged into their most similar neighbor
based on odd/even labels; if both neighbors are equally similar, they are merged into the shorter neighbor.
This yields a single sequence of non-overlapping segments that cover the entire chunk.
Each row shows the majority label on the odd and even tracks within that merged segment.
#
Start (s)
End (s)
Duration (s)
Odd label
Even label
Audio
""")
for i, seg in enumerate(ch["merged_segments"], start=1):
audio_cell = ""
if seg["audio_b64"] and seg["audio_mime"]:
audio_cell = (
''
)
html.append(
f"
{i}
"
f"
{seg['start']:.2f}
"
f"
{seg['end']:.2f}
"
f"
{seg['dur']:.2f}
"
f"
{seg['odd_label']}
"
f"
{seg['even_label']}
"
f"
{audio_cell}
"
)
html.append("
")
html.append("")
out_path = out_dir / f"seg_infer_smooth_{ts}.html"
out_path.write_text("\n".join(html), encoding="utf-8")
return out_path
# =========================
# ========= MAIN ==========
# =========================
def main():
setup_dirs()
global AUDIO_INPUT_DIR
if len(sys.argv) > 1:
AUDIO_INPUT_DIR = Path(sys.argv[1])
if not AUDIO_INPUT_DIR.is_dir():
print(f"[ERR] AUDIO_INPUT_DIR not found or not a dir: {AUDIO_INPUT_DIR}", file=sys.stderr)
sys.exit(1)
if not CKPT_PATH.is_file():
print(f"[ERR] Checkpoint not found: {CKPT_PATH}", file=sys.stderr)
sys.exit(1)
print(f"[cfg] AUDIO_INPUT_DIR = {AUDIO_INPUT_DIR}")
print(f"[cfg] OUT_DIR = {OUT_DIR}")
print(f"[cfg] CKPT_PATH = {CKPT_PATH}")
print(f"[cfg] HF_MODEL_ID = {HF_MODEL_ID}")
print(f"[cfg] ffmpeg available: {FFMPEG_AVAILABLE}")
print(f"[cfg] MIN_SEGMENT_SEC = {MIN_SEGMENT_SEC:.2f} (frames >= {MIN_SEGMENT_FRAMES})")
# find audio files
exts = {".wav", ".mp3", ".m4a", ".flac", ".ogg"}
audio_files: List[Path] = []
for p in AUDIO_INPUT_DIR.rglob("*"):
if p.is_file() and p.suffix.lower() in exts:
audio_files.append(p)
audio_files = sorted(audio_files)
if not audio_files:
print("[ERR] No audio files found.", file=sys.stderr)
sys.exit(1)
print(f"[scan] Found {len(audio_files)} audio files.")
# feature extractor
resolved, is_local = _model_resolved_name(HF_MODEL_ID)
fe = WhisperFeatureExtractor.from_pretrained(resolved, local_files_only=is_local)
# model + checkpoint
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = WhisperOddEven(HF_MODEL_ID, freeze_encoder=False).to(device)
state = torch.load(CKPT_PATH, map_location="cpu")
# accept full trainer_state dict or plain state_dict
if isinstance(state, dict) and "model" in state and any(
k.startswith("whisper.") for k in state["model"].keys()
):
state = state["model"]
missing, unexpected = model.load_state_dict(state, strict=False)
print(f"[ckpt] Loaded checkpoint from {CKPT_PATH}")
if missing:
print(f"[ckpt] Missing keys: {missing}")
if unexpected:
print(f"[ckpt] Unexpected keys: {unexpected}")
model.eval()
use_dtype = preferred_dtype()
amp_enabled = use_dtype in (torch.float16, torch.bfloat16)
chunk_results: List[Dict[str, Any]] = []
with torch.no_grad():
for fpath in audio_files:
print(f"[file] {fpath}")
try:
wav = load_audio_mono_16k(fpath)
except Exception as e:
print(f"[file] Failed to load {fpath}: {e}")
continue
chunks = split_into_chunks(wav, SAMPLE_RATE, CLIP_SECONDS)
if not chunks:
print(f"[file] No audio samples in {fpath}")
continue
for chunk_idx, start_sample, seg in chunks:
chunk_offset_sec = start_sample / SAMPLE_RATE
# features
feat = fe(seg, sampling_rate=SAMPLE_RATE, return_tensors="pt")
x = feat.input_features.to(device)
# forward
with torch.autocast(
device_type="cuda" if torch.cuda.is_available() else "cpu",
enabled=amp_enabled,
dtype=use_dtype,
):
logits = model(x)
# raw argmax
raw_ids = logits.argmax(dim=-1).squeeze(0).cpu().numpy() # [2,1500]
# aggressive smoothing with min duration per track
sm_ids = np.zeros_like(raw_ids)
for tr in range(NUM_TRACKS):
sm_ids[tr] = smooth_min_duration(raw_ids[tr], MIN_SEGMENT_FRAMES)
sm_ids_t = torch.from_numpy(sm_ids)
png_tracks = base64.b64encode(
_plot_tracks_seconds(
sm_ids_t,
f"Smoothed tracks - {fpath.name} - chunk {chunk_idx}",
)
).decode("ascii")
# merged timeline with its own min-duration smoothing
merged = build_merged_segments(sm_ids[0], sm_ids[1], MIN_SEGMENT_FRAMES)
merged_index = np.zeros(NUM_FRAMES, dtype=np.int64)
for idx, (fs, fe_, _ol, _el) in enumerate(merged, start=1):
merged_index[fs:fe_] = idx
png_merged = base64.b64encode(
_plot_merged_segments(
merged_index,
f"Merged segments - {fpath.name} - chunk {chunk_idx}",
)
).decode("ascii")
# per-track segments -> audio snippets
track_segments: List[Dict[str, Any]] = []
for tr, track_name in enumerate(("odd", "even")):
seg_runs = extract_segments(sm_ids[tr], include_bg=False)
for (lab, fs, fe_) in seg_runs:
start_t, end_t = frames_to_times(fs, fe_)
dur = end_t - start_t
if dur <= 0:
continue
sub_wav = cut_wav(seg, start_t, end_t)
if sub_wav.size == 0:
continue
try:
audio_bytes, audio_mime = wav_chunk_to_audio_bytes(sub_wav, SAMPLE_RATE)
audio_b64 = base64.b64encode(audio_bytes).decode("ascii")
except Exception as e:
print(f"[audio] Failed per-track snippet for {fpath} chunk {chunk_idx}: {e}")
audio_b64 = None
audio_mime = None
track_segments.append(
{
"track": track_name,
"label": int(lab),
"start": float(start_t),
"end": float(end_t),
"dur": float(dur),
"audio_b64": audio_b64,
"audio_mime": audio_mime,
}
)
# merged segments -> audio snippets
merged_segments: List[Dict[str, Any]] = []
for idx, (fs, fe_, odd_label, even_label) in enumerate(merged, start=1):
start_t, end_t = frames_to_times(fs, fe_)
dur = end_t - start_t
if dur <= 0:
continue
sub_wav = cut_wav(seg, start_t, end_t)
if sub_wav.size == 0:
continue
try:
audio_bytes, audio_mime = wav_chunk_to_audio_bytes(sub_wav, SAMPLE_RATE)
audio_b64 = base64.b64encode(audio_bytes).decode("ascii")
except Exception as e:
print(f"[audio] Failed merged snippet for {fpath} chunk {chunk_idx}: {e}")
audio_b64 = None
audio_mime = None
merged_segments.append(
{
"idx": idx,
"start": float(start_t),
"end": float(end_t),
"dur": float(dur),
"odd_label": int(odd_label),
"even_label": int(even_label),
"audio_b64": audio_b64,
"audio_mime": audio_mime,
}
)
chunk_results.append(
{
"file_name": fpath.name,
"chunk_idx": int(chunk_idx),
"chunk_offset": float(chunk_offset_sec),
"png_tracks": png_tracks,
"png_merged": png_merged,
"track_segments": track_segments,
"merged_segments": merged_segments,
}
)
if not chunk_results:
print("[ERR] No chunk results; nothing to write.", file=sys.stderr)
sys.exit(1)
out_html = write_html_report(OUT_DIR, chunk_results)
print(f"[done] Wrote HTML report: {out_html}")
if __name__ == "__main__":
main()
