detect_all.py

# -*- coding: utf-8 -*-
import os
import sys
import csv
import json
import glob
import argparse
import math
import warnings
import threading
import multiprocessing as mp
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import List, Tuple, Dict, Iterable

import numpy as np
from tqdm import tqdm

os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
os.environ.setdefault("FFMPEG_LOGLEVEL", "error")
os.environ.setdefault("ABSL_LOGGING_MIN_LOG_LEVEL", "3")
os.environ.setdefault("GLOG_minloglevel", "3")

warnings.filterwarnings("ignore")

# ---------------------------
# Common Utilities
# ---------------------------

SPEECH_LABELS = [
    "speech", "male speech, man speaking", "female speech, woman speaking",
    "child speech, kid speaking", "conversation", "narration, monologue",
    "babbling", "speech synthesizer", "shout", "bellow", "whoop", "yell",
    "battle cry", "children shouting", "screaming", "whispering",
    "laughter", "baby laughter", "giggle", "snicker", "belly laugh",
    "chuckle, chortle", "crying, sobbing", "baby cry, infant cry",
    "whimper", "wail, moan", "sigh", "hubbub, speech noise, speech babble",
    "chatter"
]

MUSIC_LABELS = [
    "singing", "male singing", "female singing", "choir", "yodeling",
    "chant", "mantra", "child singing", "synthetic singing", "rapping",
    "humming", "whistling", "beatboxing", "music", "musical instrument",
    "plucked string instrument", "guitar", "electric guitar", "bass guitar",
    "acoustic guitar", "steel guitar, slide guitar", "tapping (guitar technique)",
    "strum", "banjo", "sitar", "mandolin", "zither", "ukulele",
    "keyboard (musical)", "piano", "electric piano", "organ",
    "electronic organ", "hammond organ", "synthesizer", "sampler",
    "harpsichord", "percussion", "drum kit", "drum machine", "drum",
    "snare drum", "rimshot", "drum roll", "bass drum", "timpani", "tabla",
    "cymbal", "hi-hat", "wood block", "tambourine", "rattle (instrument)",
    "maraca", "gong", "tubular bells", "mallet percussion",
    "marimba, xylophone", "glockenspiel", "vibraphone", "steelpan",
    "orchestra", "brass instrument", "french horn", "trumpet", "trombone",
    "bowed string instrument", "string section", "violin, fiddle",
    "pizzicato", "cello", "double bass", "wind instrument, woodwind instrument",
    "flute", "saxophone", "clarinet", "harp", "bell", "church bell",
    "jingle bell", "bicycle bell", "tuning fork", "chime", "wind chime",
    "change ringing (campanology)", "harmonica", "accordion", "bagpipes",
    "didgeridoo", "shofar", "theremin", "singing bowl",
    "scratching (performance technique)", "pop music", "hip hop music",
    "rock music", "heavy metal", "punk rock", "grunge", "progressive rock",
    "rock and roll", "psychedelic rock", "rhythm and blues", "soul music",
    "reggae", "country", "swing music", "bluegrass", "funk", "folk music",
    "middle eastern music", "jazz", "disco", "classical music", "opera",
    "electronic music", "house music", "techno", "dubstep", "drum and bass",
    "electronica", "electronic dance music", "ambient music", "trance music",
    "music of latin america", "salsa music", "flamenco", "blues",
    "music for children", "new-age music", "vocal music", "a capella",
    "music of africa", "afrobeat", "christian music", "gospel music",
    "music of asia", "carnatic music", "music of bollywood", "ska",
    "traditional music", "independent music", "song", "background music",
    "theme music", "jingle (music)", "soundtrack music", "lullaby",
    "video game music", "christmas music", "dance music", "wedding music",
    "happy music", "funny music", "sad music", "tender music",
    "exciting music", "angry music", "scary music"
]

BIRD_LABELS = [
    "bird", "bird vocalization, bird call, bird song", "chirp, tweet",
    "squawk", "pigeon, dove", "coo", "crow", "caw", "owl", "hoot",
    "bird flight, flapping wings", "fowl", "chicken, rooster", "cluck",
    "crowing, cock-a-doodle-doo", "turkey", "gobble", "duck", "quack",
    "goose", "honk",
]

WATER_LABELS = [
    "water", "rain", "raindrop", "rain on surface", "stream", "waterfall",
    "ocean", "waves, surf", "steam", "gurgling", "water tap, faucet",
    "sink (filling or washing)", "bathtub (filling or washing)", "toilet flush",
    "liquid", "splash, splatter", "slosh", "squish", "drip", "pour",
    "trickle, dribble", "gush", "fill (with liquid)", "spray",
    "pump (liquid)", "stir", "boiling"
]

WIND_LABELS = [
    "wind", "wind noise (microphone)", "rustling leaves",
]

ENGINE_LABELS = [
    "aircraft engine", "jet engine", "engine", "light engine (high frequency)",
    "medium engine (mid frequency)", "heavy engine (low frequency)",
    "engine knocking", "engine starting", "idling", "accelerating, revving, vroom"
]

INSECT_LABELS = [
    "insect", "cricket", "mosquito", "fly, housefly", "buzz", "bee, wasp, etc.",
]

# Unified canonical labels for fusion/metrics
CANON_LABELS = ["speech", "music", "bird", "water", "wind", "engine", "insect"]
LABEL_MAP = {
    "speech": SPEECH_LABELS,
    "music": MUSIC_LABELS,
    "bird":  BIRD_LABELS,
    "water": WATER_LABELS,
    "wind": WIND_LABELS,
    "engine": ENGINE_LABELS,
    "insect": INSECT_LABELS,
}

def _list_audio(work: Path) -> List[str]:
    return sorted(glob.glob(str(work/"*.wav")) + glob.glob(str(work/"*.flac")))

def _write_csv(out_csv: Path, segments: Dict[str, List[Tuple[str, float, float]]]):
    out_csv.parent.mkdir(parents=True, exist_ok=True)
    with open(out_csv, "w", newline="", encoding="utf-8") as f:
        wr = csv.writer(f)
        wr.writerow(["filename", "start", "end", "label"])
        for fp, segs in segments.items():
            for lab, s, e in segs:
                wr.writerow([Path(fp).name, s, e, lab])

def _read_csv_segments(out_csv: Path, files: List[str]) -> Dict[str, List[Tuple[str, float, float]]]:
    name2full = {Path(fp).name: fp for fp in files}
    segments: Dict[str, List[Tuple[str, float, float]]] = {}
    
    with open(out_csv, "r", newline="", encoding="utf-8") as f:
        rd = csv.DictReader(f)
        for row in rd:
            fn = row["filename"]
            fp = name2full.get(fn)
            if fp is None:
                continue  # Skip if present in CSV but missing in file list

            s = float(row["start"])
            e = float(row["end"])
            lab = row["label"]
            segments.setdefault(fp, []).append((lab, s, e))

    return segments

def _union_duration(segs, ih_labels=("speech", "music")):
    ih_labels = set(ih_labels)
    ivals = [(s, e) for lab, s, e in segs if lab in ih_labels and e > s]
    if not ivals:
        return 0.0
    ivals.sort(key=lambda x: x[0])
    merged = []
    cs, ce = ivals[0]
    for s, e in ivals[1:]:
        if s <= ce:
            ce = max(ce, e)
        else:
            merged.append((cs, ce))
            cs, ce = s, e
    merged.append((cs, ce))
    return sum(e - s for s, e in merged)

def _calc_ih(files: List[str], segments: Dict[str, List[Tuple[str, float, float]]], duration: float, ih_labels=("speech","music")):
    metrics = {}
    ih_count = 0
    ih_dur_sum = 0.0
    total_dur_sum = 0.0
    for fp in files:
        segs = segments.get(fp, [])
        dur = duration
        # Use union duration
        ih = _union_duration(segs, ih_labels)
        has_ih = ih > 0.0
        metrics[Path(fp).name] = {
            "IH@vid": int(has_ih),
            "IH-dur": (ih / dur) if dur else 0.0
        }
        if has_ih:
            ih_count += 1
        ih_dur_sum += ih
        total_dur_sum += dur
        
    overall = {
        "IH@vid": (ih_count / len(files)) if files else 0.0,
        "IH-dur": (ih_dur_sum / total_dur_sum) if total_dur_sum else 0.0
    }
    aggregates = {
        "num_files": len(files),
        "ih_video_count": ih_count,
        "ih_duration_sum": ih_dur_sum,
        "total_duration": total_dur_sum
    }
    return overall, metrics, aggregates

def _labels_suffix(labels: Iterable[str]) -> str:
    """Generate a label suffix for filenames, e.g., ('speech','music') -> 'speech_music'"""
    labs = [str(x).strip().lower() for x in labels if str(x).strip()]
    labs = sorted(set(labs))
    return "_".join(labs) if labs else "none"

def _save_json(path: Path, payload: dict):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=4, ensure_ascii=False)

def _load_json(path: Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

# ---------------------------
# Post-processing: Time Grid & Merging
# ---------------------------

def _merge_frames(frames: List[Tuple[float, float, str]],
                  min_dur: float = 0.20,
                  min_gap: float = 0.15) -> List[Tuple[str, float, float]]:
    if not frames:
        return []
    frames = sorted(frames, key=lambda x: x[0])
    out = []
    cur_l, cur_s, cur_e = frames[0][2], frames[0][0], frames[0][1]
    for s, e, l in frames[1:]:
        if l == cur_l and s - cur_e <= min_gap:
            cur_e = e
        else:
            if cur_e - cur_s >= min_dur:
                out.append((cur_l, round(cur_s, 3), round(cur_e, 3)))
            cur_l, cur_s, cur_e = l, s, e
    if cur_e - cur_s >= min_dur:
        out.append((cur_l, round(cur_s, 3), round(cur_e, 3)))
    return out

def _fuse_segments(segment_lists, rule="mv",
                   labels=CANON_LABELS,
                   hop=0.02, min_dur=0.20, min_gap=0.15,
                   tolerance=0.05):
    labels = ["speech", "music"]
    T = max((e for segs in segment_lists for _, _, e in segs), default=0.0)
    if T == 0: return []
    n = int(math.ceil(T / hop))
    frames = []
    for lab in labels:
        votes = np.zeros((len(segment_lists), n), dtype=np.uint8)
        for k, segs in enumerate(segment_lists):
            for l, s, e in segs:
                if l != lab: continue
                i0 = max(0, int(math.floor((s - tolerance) / hop)))
                i1 = min(n, int(math.ceil((e + tolerance) / hop)))
                votes[k, i0:i1] = 1
                
        # Voting rules
        if rule == "or":
            active = (votes.sum(axis=0) >= 1)
        elif rule == "and":
            active = (votes.sum(axis=0) == votes.shape[0])
        else:  # majority vote (mv)
            active = (votes.sum(axis=0) >= (votes.shape[0] + 1)//2)
            
        # Merge consecutive frames
        i = 0
        while i < n:
            if active[i]:
                j = i+1
                while j < n and active[j]:
                    j += 1
                frames.append((i*hop, j*hop, lab))
                i = j
            else:
                i += 1
    return _merge_frames(frames, min_dur=min_dur, min_gap=min_gap)

# ---------------------------
# Detectors
# ---------------------------

class BaseDetector:
    name = "base"
    def segment(self, fpath: str) -> List[Tuple[str, float, float]]:
        raise NotImplementedError

# 1) inaSpeechSegmenter
_tls = threading.local()
def _get_inaseg():
    if not hasattr(_tls, "seg"):
        from inaSpeechSegmenter import Segmenter
        _tls.seg = Segmenter(vad_engine='smn', detect_gender=False)
    return _tls.seg

# Global singleton for inaSpeechSegmenter in subprocesses
_INA_SEG_SINGLETON = None

def _init_ina_in_this_process():
    """Initialize inaSpeechSegmenter once per subprocess"""
    global _INA_SEG_SINGLETON
    if _INA_SEG_SINGLETON is not None:
        return _INA_SEG_SINGLETON

    os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # Disable GPU
    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")
    os.environ.setdefault("TF_NUM_INTRAOP_THREADS", "1")
    os.environ.setdefault("TF_NUM_INTEROP_THREADS", "1")

    from inaSpeechSegmenter import Segmenter
    _INA_SEG_SINGLETON = Segmenter(vad_engine='smn', detect_gender=False)
    return _INA_SEG_SINGLETON

def _ina_worker(fp: str):
    try:
        seg = _init_ina_in_this_process()
        raw = [(lab, float(s), float(e)) for lab, s, e in seg(fp)]
        keep = [(lab, s, e) for lab, s, e in raw if lab in ("speech", "music")]
        segs = [(lab, round(s,3), round(e,3)) for lab, s, e in keep]
        return fp, segs, None
    except Exception as e:
        return fp, [], str(e)


class InaDetector(BaseDetector):
    name = "ina"
    def segment(self, fpath: str):
        seg = _get_inaseg()
        raw = [(lab, float(s), float(e)) for lab, s, e in seg(fpath)]
        keep = [(lab, s, e) for lab, s, e in raw if lab in ("speech", "music")]
        return [(lab, round(s,3), round(e,3)) for lab, s, e in keep]


# ========= YAMNet: Subprocess singleton + worker (CPU-only) =========
_YAMNET_SINGLETON = None
_YAMNET_IDX = None
_YAMNET_CONF = {"thr": {"speech":0.40, "music":0.30, "bird":0.30, "water":0.30}, "hop": 0.48}

def _init_yamnet_in_this_process():
    """
    Initialize YAMNet (CPU-only) once per process.
    Returns (model, librosa, idx_speech, idx_music)
    """
    global _YAMNET_SINGLETON, _YAMNET_IDX

    if _YAMNET_SINGLETON is not None and _YAMNET_IDX is not None:
        return _YAMNET_SINGLETON, _YAMNET_IDX

    # Disable GPU and limit threads before importing TF
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")
    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")
    os.environ.setdefault("TF_NUM_INTRAOP_THREADS", "1")
    os.environ.setdefault("TF_NUM_INTEROP_THREADS", "1")

    import tensorflow as tf
    import tensorflow_hub as hub
    import librosa
    try:
        tf.config.set_visible_devices([], 'GPU')
    except Exception:
        pass
    tf.config.threading.set_intra_op_parallelism_threads(1)
    tf.config.threading.set_inter_op_parallelism_threads(1)

    model = hub.load("https://tfhub.dev/google/yamnet/1")
    class_map = model.class_map_path().numpy()
    with tf.io.gfile.GFile(class_map) as f:
        import csv as _csv
        reader = _csv.DictReader(f)
        names = [r["display_name"].strip().lower() for r in reader]
    
    label_sets = {k: {x.strip().lower() for x in v} for k, v in LABEL_MAP.items()}
    idx = {k: [i for i, n in enumerate(names) if n in s] for k, s in label_sets.items()}

    _YAMNET_SINGLETON = (model, librosa)
    _YAMNET_IDX = idx
    return _YAMNET_SINGLETON, _YAMNET_IDX

def _yamnet_worker(fp: str):
    try:
        (model, librosa), idx = _init_yamnet_in_this_process()
        hop = _YAMNET_CONF["hop"]
        thr = _YAMNET_CONF["thr"]

        y, sr = librosa.load(fp, sr=16000, mono=True)
        scores, embeddings, spectrogram = model(y)
        scores = scores.numpy()

        frames = []
        for i in range(scores.shape[0]):
            s = i * hop
            e = (i + 1) * hop
            for lab in CANON_LABELS:
                ids = idx.get(lab, [])
                p = scores[i, ids].max() if ids else 0.0
                if p >= thr.get(lab, 0.30):
                    frames.append((s, e, lab))

        segs = _merge_frames(frames, min_dur=0.20, min_gap=0.15)
        segs = [(lab, float(s), float(e)) for (lab, s, e) in segs]
        return fp, segs, None
    except Exception as e:
        return fp, [], str(e)


# 2) YAMNet
class YamnetDetector(BaseDetector):
    name = "yamnet"
    def __init__(self, thr=None, hop=0.48):
        os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
        import tensorflow as tf
        import tensorflow_hub as hub
        import librosa
        self.tf = tf
        self.hub = hub
        self.librosa = librosa
        self.model = hub.load("https://tfhub.dev/google/yamnet/1")

        class_map = self.model.class_map_path().numpy()
        with tf.io.gfile.GFile(class_map) as f:
            import csv as _csv
            reader = _csv.DictReader(f)
            names = [r["display_name"].strip().lower() for r in reader]

        label_sets = {k: {x.strip().lower() for x in v} for k, v in LABEL_MAP.items()}
        self.idx = {k: [i for i, n in enumerate(names) if n in s] for k, s in label_sets.items()}

        self.thr = thr or {"speech":0.40, "music":0.30, "bird":0.30, "water":0.30}
        self.hop = hop

    def segment(self, fpath: str):
        y, sr = self.librosa.load(fpath, sr=16000, mono=True)
        scores, embeddings, spectrogram = self.model(y)
        scores = scores.numpy()

        frames = []
        for i in range(scores.shape[0]):
            s = i * self.hop
            e = (i + 1) * self.hop
            for lab in CANON_LABELS:
                ids = self.idx.get(lab, [])
                p = scores[i, ids].max() if ids else 0.0
                if p >= self.thr.get(lab, 0.30):
                    frames.append((s, e, lab))

        return _merge_frames(frames, min_dur=0.20, min_gap=0.15)


# 3) PANNs CNN14 (Frame-level)
class PannsDetector(BaseDetector):
    name = "panns"
    def __init__(self, thr_speech: float = 0.4, thr_music: float = 0.30, device: str = "cuda"):
        import librosa
        import numpy as np
        from panns_inference import SoundEventDetection, labels as panns_labels
        self.librosa = librosa
        self.np = np
        self.sr = 32000
        self.model = SoundEventDetection(checkpoint_path=None, device=device)
        self.labels = [str(x).strip().lower() for x in panns_labels]
        self.idx_speech = [i for i, n in enumerate(self.labels) if n in SPEECH_LABELS]
        self.idx_music  = [i for i, n in enumerate(self.labels) if n in MUSIC_LABELS]
        self.thr_speech = float(thr_speech)
        self.thr_music = float(thr_music)

        label_sets = {k: {x.strip().lower() for x in v} for k, v in LABEL_MAP.items()}
        self.idx = {k: [i for i, n in enumerate(self.labels) if n in s] for k, s in label_sets.items()}
        self.thr = {"speech":0.40, "music":0.30, "bird":0.30, "water":0.30}

    def segment(self, fpath: str):
        y, _ = self.librosa.load(fpath, sr=self.sr, mono=True)
        if len(y) == 0:
            return []
        out = self.model.inference(y[None, :])
        if isinstance(out, dict) and "framewise_output" in out:
            framewise = out["framewise_output"]
        else:
            framewise = out
        framewise = self.np.asarray(framewise)
        if framewise.ndim == 3 and framewise.shape[0] == 1:
            framewise = framewise[0]
        if framewise.ndim != 2 or framewise.shape[0] == 0:
            return []
        T = framewise.shape[0]
        duration = len(y) / self.sr
        t_hop = duration / T
        frames = []
        for i in range(T):
            s = i * t_hop
            e = (i + 1) * t_hop
            for lab in CANON_LABELS:
                ids = self.idx.get(lab, [])
                p = float(framewise[i, ids].max()) if ids else 0.0
                if p >= self.thr.get(lab, 0.40):
                    frames.append((s, e, lab))
        return _merge_frames(frames, min_dur=0.20, min_gap=0.15)

# ---------------------------
# Detect: Process single directory
# ---------------------------
def run_detect(work_dir: str, detector: str, jobs: int = 6, duration: float = 9.98, ih_labels=("speech","music"), skip_existing: bool = True) -> Path:
    work = Path(work_dir)
    parents = work.parents
    l2 = parents[0].name if len(parents) >= 1 else ""
    l1 = parents[1].name if len(parents) >= 2 else ""
    files = _list_audio(work)
    cache = work / "cache"
    cache.mkdir(exist_ok=True)

    det = detector.lower()
    if det == "ina":
        _init_ina_in_this_process()
        DET = InaDetector()
    elif det == "yamnet":
        DET = YamnetDetector()
    elif det == "panns":
        DET = PannsDetector()
    else:
        raise ValueError(f"Unknown detector: {det}")

    out_csv  = cache / f"segment_{det}.csv"
    ih_suf = _labels_suffix(ih_labels)
    out_json = cache / f"ih_metrics_{det}_{ih_suf}.json"

    # Skip existing cache
    if skip_existing and out_csv.exists() and out_json.exists():
        print(f"[{det}] Found cache, skipping computation -> {out_json}")
        return out_json

    # Recover segments from CSV if JSON is missing
    if skip_existing and out_csv.exists() and (not out_json.exists()):
        print(f"[{det}] Found CSV but missing JSON, recovering from CSV -> {out_json}")
        segments = _read_csv_segments(out_csv, work_dir)
        overall, metrics, aggregates = _calc_ih(files, segments, duration, ih_labels=ih_labels)
        payload = {
            "Detector": det,
            "IH Labels": list(ih_labels),
            "Overall Metrics": overall,
            "Aggregates": aggregates,
            "Individual Audio Metrics": metrics,
            "Segments": {Path(fp).name: segs for fp, segs in segments.items()},
            "Errors": []
        }
        _save_json(out_json, payload)
        print(f"[{det}] CSV : {out_csv}")
        print(f"[{det}] JSON: {out_json}")
        return out_json

    segments: Dict[str, List[Tuple[str, float, float]]] = {}
    errors = []

    def _process_one(fp: str):
        try:
            segs = DET.segment(fp)
            return fp, segs, None
        except Exception as e:
            return fp, [], str(e)

    use_mp = (max(1, jobs) > 1) and (det in ("ina", "yamnet"))
    if use_mp:
        try:
            mp.set_start_method("spawn", force=True)  # Prevent fork+TF issues
        except RuntimeError:
            pass

        with ProcessPoolExecutor(max_workers=max(1, jobs)) as ex:
            if det == "ina":
                futs = [ex.submit(_ina_worker, fp) for fp in files]
            elif det == "yamnet":
                # Pass thresholds/hop config to worker
                _YAMNET_CONF.update({"thr_speech": 0.40, "thr_music": 0.30, "hop": 0.48})
                futs = [ex.submit(_yamnet_worker, fp) for fp in files]
                
            for fut in tqdm(as_completed(futs), total=len(futs), desc=f"[{det}: {l1}/{l2}]", unit="file"):
                fp, segs, err = fut.result()
                if err:
                    errors.append({"file": Path(fp).name, "error": str(err)})
                else:
                    segments[fp] = segs
    else:
        # Fallback to single process
        for fp in tqdm(files, desc=f"[{det}: {l1}/{l2}]", unit="file"):
            try:
                segs = DET.segment(fp)
                segments[fp] = segs
            except Exception as e:
                errors.append({"file": Path(fp).name, "error": str(e)})

    _write_csv(out_csv, segments)
    overall, metrics, aggregates = _calc_ih(files, segments, duration, ih_labels=ih_labels)
    payload = {
        "Detector": det,
        "IH Labels": list(ih_labels),
        "Overall Metrics": overall,
        "Aggregates": aggregates,
        "Individual Audio Metrics": metrics,
        "Segments": {Path(fp).name: segs for fp, segs in segments.items()},
        "Errors": errors
    }
    _save_json(out_json, payload)
    print(f"[{det}] CSV : {out_csv}")
    print(f"[{det}] JSON: {out_json}")
    return out_json

# ---------------------------
# Fuse: Process single directory
# ---------------------------
def run_fuse(work_dir: str, hop=0.02, min_dur=0.20, min_gap=0.15, duration: float = 9.98,
             ih_labels=("speech","music"), skip_existing: bool = True) -> Path:
    work = Path(work_dir)
    files = _list_audio(work)
    cache = work / "cache"
    cache.mkdir(exist_ok=True)

    ih_suf = _labels_suffix(ih_labels)
    out_json_all = cache / f"ih_metrics_fused_all_{ih_suf}.json"
    out_csv_and = cache / "segment_fused_and.csv"
    out_csv_mv  = cache / "segment_fused_mv.csv"
    out_csv_or  = cache / "segment_fused_or.csv"

    # Skip existing cache
    if skip_existing and out_json_all.exists() and out_csv_and.exists() and out_csv_mv.exists() and out_csv_or.exists():
        print(f"[FUSE] Found cache, skipping fusion -> {out_json_all}")
        return out_json_all

    cands = sorted(cache.glob("ih_metrics_*.json"))
    input_jsons = [
        p for p in cands
        if not p.name.startswith("ih_metrics_fused")  # Exclude previously fused outputs
    ]
    if not input_jsons:
        print(f"[FUSE] No input JSONs found (expected {cache}/ih_metrics_*.json)")
        return out_json_all

    print(f"[FUSE] ({work}) Found {len(input_jsons)} input JSONs:")
    for p in input_jsons:
        print("   ", p.name)

    banks: List[Dict[str, List[Tuple[str, float, float]]]] = []
    for jpath in input_jsons:
        obj = _load_json(jpath)
        segs = obj.get("Segments", {})
        bank = {}
        for fp in files:
            name = Path(fp).name
            bank[fp] = [(lab, float(s), float(e)) for lab, s, e in segs.get(name, [])]
        banks.append(bank)

    def _do_one_fusion(rule: str):
        fused: Dict[str, List[Tuple[str, float, float]]] = {}
        for fp in files:
            cand = [b[fp] for b in banks]
            fused[fp] = _fuse_segments(cand, rule=rule, labels=CANON_LABELS, hop=hop, min_dur=min_dur, min_gap=min_gap)
        return fused

    rules = ["and", "mv", "or"]
    fused_map: Dict[str, Dict[str, List[Tuple[str, float, float]]]] = {}
    overall_map: Dict[str, Dict[str, float]] = {}
    metrics_map: Dict[str, Dict[str, Dict[str, float]]] = {}
    aggregates_map: Dict[str, Dict[str, float]] = {}

    for r in rules:
        fused_r = _do_one_fusion(r)
        fused_map[r] = fused_r
        out_csv_r = cache / f"segment_fused_{r}.csv"
        _write_csv(out_csv_r, fused_r)
        print(f"[FUSE-{r.upper()}] CSV : {out_csv_r}")
        
        overall_r, metrics_r, aggregates_r = _calc_ih(files, fused_r, duration, ih_labels=ih_labels)
        overall_map[r.upper()] = overall_r
        metrics_map[r.upper()] = metrics_r
        aggregates_map[r.upper()] = aggregates_r

    payload_all = {
        "Fusions": [x.upper() for x in rules],
        "Inputs": [p.name for p in input_jsons],
        "IH Labels": list(ih_labels),
        "Overall Metrics": overall_map,
        "Aggregates": aggregates_map,
        "Individual Audio Metrics": metrics_map,
        "Segments": {
            r.upper(): {Path(fp).name: segs for fp, segs in fused_map[r].items()}
            for r in rules
        }
    }
    _save_json(out_json_all, payload_all)
    print(f"[FUSE-ALL] JSON: {out_json_all}")
    return out_json_all

# ---------------------------
# Directory Traversal
# ---------------------------
def iter_work_dirs_old(base_root: Path, sub: str) -> Iterable[Path]:
    """Traverse base_root/L1/L2/sub and yield existing work_dir Paths"""
    for l1 in sorted(p for p in base_root.iterdir() if p.is_dir()):
        if l1.name in ['.cache', 'Music', 'Others']:
            print(f"Skipping directory: {l1.name}")
            continue
        for l2 in sorted(p for p in l1.iterdir() if p.is_dir()):
            if l2.name in ['Human voice', 'Whistling', 'Other sourceless', 'Specific impact sounds']:
                print(f"Skipping directory: {l2.name}")
                continue
            wd = l2 / sub
            if wd.exists() and wd.is_dir():
                yield wd

def iter_work_dirs(base_root: Path, sub: str) -> Iterable[Path]:
    """Traverse base_root/L1/L2/sub and yield existing work_dir Paths"""
    for l1 in sorted(p for p in base_root.iterdir() if p.is_dir()):
        for l2 in sorted(p for p in l1.iterdir() if p.is_dir()):
            wd = l2 / sub
            if wd.exists() and wd.is_dir():
                yield wd

# ---------------------------
# Accumulation & Summary
# ---------------------------
def _count_weight_accumulate(overall: dict, n_files: int, acc: dict):
    if n_files <= 0 or not overall:
        return
    acc.setdefault("weight", 0)
    acc.setdefault("ih_vid_sum", 0.0)
    acc.setdefault("ih_dur_sum", 0.0)
    acc["weight"] += n_files
    acc["ih_vid_sum"] += float(overall.get("IH@vid", 0.0)) * n_files
    acc["ih_dur_sum"] += float(overall.get("IH-dur", 0.0)) * n_files

def _count_weight_finalize(acc: dict) -> dict:
    w = acc.get("weight", 0)
    if w == 0:
        return {"IH@vid": 0.0, "IH-dur": 0.0, "file_count": 0}
    return {
        "IH@vid": acc["ih_vid_sum"] / w,
        "IH-dur": acc["ih_dur_sum"] / w,
        "file_count": w
    }

# ---------------------------
# Batch Detect All Subdirectories
# ---------------------------
def run_detect_all(base_root: str, sub: str, detector: str, jobs: int = 4, duration: float = 9.98, ih_labels=("speech","music"), skip_existing: bool = True):
    base = Path(base_root)
    if not base.exists():
        print(f"[ERROR] base_root does not exist: {base}", file=sys.stderr)
        sys.exit(1)

    l1_acc: Dict[str, dict] = {}
    l2_metrics: Dict[str, dict] = {}
    overall_acc: dict = {}

    total = 0
    for work_dir in iter_work_dirs(base, sub):
        print(f"\n[DETECT] work_dir = {work_dir}")
        json_path = run_detect(str(work_dir), detector=detector, jobs=jobs, duration=duration, ih_labels=ih_labels, skip_existing=skip_existing)
        if not Path(json_path).exists():
            continue
        obj = _load_json(json_path)
        overall = obj.get("Overall Metrics") or {}
        indiv = obj.get("Individual Audio Metrics") or {}
        n_files = len(indiv)

        if sub:
            l2_key = str(work_dir.parent)          # L1/L2
            l1_key = str(work_dir.parent.parent)   # L1
        else:
            l2_key = str(work_dir.stem)            # L1/L2
            l1_key = str(work_dir.parent)          # L1
            
        l2_metrics[l2_key] = overall
        l1_acc.setdefault(l1_key, {})
        _count_weight_accumulate(overall, n_files, l1_acc[l1_key])
        _count_weight_accumulate(overall, n_files, overall_acc)
        total += 1

    l1_metrics = {k: _count_weight_finalize(v) for k, v in l1_acc.items()}
    overall_metrics = _count_weight_finalize(overall_acc)

    print(f"\n[DETECT-ALL] Traversal complete. Processed {total} directories.")
    print("\n===== L1 Average (Weighted) =====")
    for k in sorted(l1_metrics.keys()):
        m = l1_metrics[k]
        print(f"{k} : IH@vid={m['IH@vid']:.4f} | IH-dur={m['IH-dur']:.4f}")
    print("\n===== GLOBAL Overall =====")
    print(f"IH@vid={overall_metrics['IH@vid']:.4f} | IH-dur={overall_metrics['IH-dur']:.4f} | files={overall_metrics['file_count']}")

    summary = {
        "Mode": "detect_all",
        "Detector": detector,
        "Sub": sub,
        "BaseRoot": base_root,
        "L1 Metrics": l1_metrics,
        "L2 Metrics": l2_metrics,
        "Global Overall Metrics": overall_metrics
    }
    ih_suf = _labels_suffix(ih_labels)
    out_summary = base / f"summary_detect_all_{ih_suf}_{detector}.json"
    _save_json(out_summary, summary)
    print(f"[DETECT-ALL] Summary: {out_summary}")

# ---------------------------
# Batch Fuse All Subdirectories
# ---------------------------
def run_fuse_all(base_root: str, sub: str, hop=0.02, min_dur=0.20, min_gap=0.15, duration: float = 9.98, ih_labels=("speech","music"), skip_existing: bool = True):
    base = Path(base_root)
    if not base.exists():
        print(f"[ERROR] base_root does not exist: {base}", file=sys.stderr)
        sys.exit(1)

    l1_acc_map = {"AND": {}, "MV": {}, "OR": {}}
    l2_metrics_map = {"AND": {}, "MV": {}, "OR": {}}
    overall_acc_map = {"AND": {}, "MV": {}, "OR": {}}

    total = 0
    for work_dir in iter_work_dirs(base, sub):
        print(f"\n[FUSE] work_dir = {work_dir}")
        json_path = run_fuse(str(work_dir), hop=hop, min_dur=min_dur, min_gap=min_gap, duration=duration, skip_existing=skip_existing)
        if not Path(json_path).exists():
            continue
        obj = _load_json(json_path)
        overall_map = obj.get("Overall Metrics") or {}
        indiv_map = obj.get("Individual Audio Metrics") or {}
        l1_key = str(work_dir.parent.parent)   # L1
        l2_key = str(work_dir.parent)          # L1/L2

        for rule in ("AND", "MV", "OR"):
            overall = overall_map.get(rule) or {}
            indiv = indiv_map.get(rule) or {}
            n_files = len(indiv)
            if overall:
                l2_metrics_map[rule][l2_key] = overall
            l1_acc_map[rule].setdefault(l1_key, {})
            _count_weight_accumulate(overall, n_files, l1_acc_map[rule][l1_key])
            _count_weight_accumulate(overall, n_files, overall_acc_map[rule])
        total += 1

    l1_metrics_map = {r: {k: _count_weight_finalize(v) for k, v in acc.items()} for r, acc in l1_acc_map.items()}
    overall_metrics = {r: _count_weight_finalize(acc) for r, acc in overall_acc_map.items()}

    print(f"\n[FUSE-ALL] Traversal complete. Processed {total} directories.")
    for rule in ("AND", "MV", "OR"):
        print(f"\n== {rule} ==")
        for k in sorted(l1_metrics_map[rule].keys()):
            m = l1_metrics_map[rule][k]
            print(f"{k} : IH@vid={m['IH@vid']:.4f} | IH-dur={m['IH-dur']:.4f}")
        gm = overall_metrics[rule]
        print(f"GLOBAL: IH@vid={gm['IH@vid']:.4f} | IH-dur={gm['IH-dur']:.4f} | files={gm['file_count']}")

    summary = {
        "Mode": "fuse_all",
        "Sub": sub,
        "BaseRoot": base_root,
        "L1 Metrics": l1_metrics_map,
        "L2 Metrics": l2_metrics_map,
        "Global Overall Metrics": overall_metrics
    }
    out_summary = base / f"summary_fuse_all_{sub}.json"
    _save_json(out_summary, summary)
    out_global = base / f"global_overall_fuse_all_{sub}.json"
    _save_json(out_global, {"Global Overall Metrics": overall_metrics})
    print(f"[FUSE-ALL] Summary: {out_summary}")
    print(f"[FUSE-ALL] Global Overall: {out_global}")


def _parse_ih_labels(s: str):
    labs = [x.strip().lower() for x in s.split(",") if x.strip()]
    labs = [x for x in labs if x in CANON_LABELS]
    
    if labs:
        return tuple(labs)
    else:
        # Raise value error with reasoning
        raise ValueError(f"No valid labels found in input: '{s}'")


MODE_CONFIG = {
    "gt": {
        "base_root": "/workspace/Kling-Audio-Eval",
        "sub": "audio"
    },
    "thinksound": {
        "base_root": "/workspace/ThinkSound",
        "sub": ""
    },
    "speech": {
        "base_root": "/workspace/speech",
        "sub": ""
    },
    "music": {
        "base_root": "/workspace/music",
        "sub": ""
    },
    "water": {
        "base_root": "/workspace/water",
        "sub": ""
    },
    "bird": {
        "base_root": "/workspace/bird",
        "sub": ""
    },
    "insect": {
        "base_root": "/workspace/insect",
        "sub": ""
    },
    "adj-music": {
        "base_root": "/workspace/adj_replace/music",
        "sub": ""
    },
    "adj-speech": {
        "base_root": "/workspace/adj_replace/speech",
        "sub": ""
    },
    "retri-music": {
        "base_root": "/workspace/retri_replace/music",
        "sub": ""
    },
    "retri-speech": {
        "base_root": "/workspace/retri_replace/speech",
        "sub": ""
    },
}

# ---------------------------
# CLI
# ---------------------------
def main():
    ap = argparse.ArgumentParser(description="IH Segmentation and Fusion (Supports L1/L2/sub traversal: detect/fuse + aggregation + global output)")
    subp = ap.add_subparsers(dest="cmd", required=True)

    # Single directory detect
    ap_det = subp.add_parser("detect", help="Run a specific detector on a single work_dir and export JSON/CSV")
    ap_det.add_argument("--work_dir", default="/home/jovyan/shared/liyangchen/v2a/AVE_Dataset/prediction/thinksound/0910_batch_size8")
    ap_det.add_argument("--detector", required=True, choices=["ina","yamnet","panns"])
    ap_det.add_argument("--jobs", type=int, default=8)
    ap_det.add_argument("--duration", type=float, default=9.98)
    ap_det.add_argument("--ih_labels", default="speech,music", help="Comma-separated, e.g., speech,music or speech,music,bird or bird,water")
    ap_det.add_argument("--no_skip_existing", action="store_true", help="Do not skip existing cache files")

    # Single directory fuse
    ap_fuse = subp.add_parser("fuse", help="Read ih_metrics_*.json from cache for a single work_dir and output AND/MV/OR fusions")
    ap_fuse.add_argument("--work_dir", default="/home/jovyan/shared/liyangchen/v2a/AVE_Dataset/prediction/thinksound/0910_batch_size8")
    ap_fuse.add_argument("--hop", type=float, default=0.01)
    ap_fuse.add_argument("--min_dur", type=float, default=0.20)
    ap_fuse.add_argument("--min_gap", type=float, default=0.15)
    ap_fuse.add_argument("--duration", type=float, default=9.98)
    ap_fuse.add_argument("--no_skip_existing", action="store_true", help="Do not skip existing cache files")

    # Detect across all directories
    ap_det_all = subp.add_parser("detect_all", help="Traverse base_root/L1/L2/sub, run detect per directory, and output L1/L2/Global weighted summaries")
    ap_det_all.add_argument("--base_root", required=True, help="Data root directory, e.g., /workspace/ThinkSound")
    ap_det_all.add_argument("--sub", default="", help="Subdirectory name, default is empty (search audio directly under L2)")
    ap_det_all.add_argument("--detector", required=True, choices=["ina","yamnet","panns"])
    ap_det_all.add_argument("--jobs", type=int, default=8)
    ap_det_all.add_argument("--duration", type=float, default=9.98)
    ap_det_all.add_argument("--ih_labels", default="speech,music", help="Comma-separated, e.g., speech,music or speech,music,bird or bird,water")
    ap_det_all.add_argument("--no_skip_existing", action="store_true")

    # Fuse across all directories
    ap_fuse_all = subp.add_parser("fuse_all", help="Traverse base_root/L1/L2/sub, run fuse per directory, and output L1/L2/Global weighted summaries")
    ap_fuse_all.add_argument("--base_root", required=True, help="Data root directory, e.g., /workspace/ThinkSound")
    ap_fuse_all.add_argument("--sub", default="", help="Subdirectory name, default is empty (search audio directly under L2)")
    ap_fuse_all.add_argument("--hop", type=float, default=0.01)
    ap_fuse_all.add_argument("--min_dur", type=float, default=0.20)
    ap_fuse_all.add_argument("--min_gap", type=float, default=0.15)
    ap_fuse_all.add_argument("--duration", type=float, default=9.98)
    ap_fuse_all.add_argument("--ih_labels", default="speech,music", help="Comma-separated, e.g., speech,music or speech,music,bird or bird,water")
    ap_fuse_all.add_argument("--no_skip_existing", action="store_true")

    args = ap.parse_args()
    ih_labels = _parse_ih_labels(getattr(args, "ih_labels", "speech,music"))

    if args.cmd == "detect":
        run_detect(args.work_dir, args.detector, args.jobs, args.duration, ih_labels=ih_labels, skip_existing=(not args.no_skip_existing))
    elif args.cmd == "fuse":
        run_fuse(args.work_dir, hop=args.hop, min_dur=args.min_dur, min_gap=args.min_gap, duration=args.duration, ih_labels=ih_labels, skip_existing=(not args.no_skip_existing))
    elif args.cmd == "detect_all":
        run_detect_all(args.base_root, args.sub, args.detector, args.jobs, args.duration, ih_labels=ih_labels, skip_existing=(not args.no_skip_existing))
    elif args.cmd == "fuse_all":
        run_fuse_all(args.base_root, args.sub, hop=args.hop, min_dur=args.min_dur, min_gap=args.min_gap, duration=args.duration, ih_labels=ih_labels, skip_existing=(not args.no_skip_existing))

if __name__ == "__main__":
    main()
    os.system("stty sane")