Update app.py

app.py

@@ -1,25 +1,26 @@
 """
 DCASE 2024 Task 4 baseline (CRNN + BEATs) demo for Hugging Face Spaces (ZeroGPU compatible).
-
-Files expected in the Space root:
-  - last.ckpt                    (DCASE 2024 baseline checkpoint)
-  - BEATs_iter3_plus_AS2M.pt     (optional here: if absent, the app tries to download it)
 """
 import os
 import sys
-import subprocess
-
-# --- CRITICAL FIX FOR MISSING DEPENDENCIES ---
-# Downgrades setuptools for dcase_util compatibility and catches missing evaluation packages
-subprocess.run([sys.executable, "-m", "pip", "install", "-q", "setuptools<70.0.0", "psds_eval", "codecarbon", "thop"], check=False)
-# ----------------------------------------------
+import types
+
+# --- CRITICAL FIX: MOCK PKG_RESOURCES ---
+# dcase_util requires pkg_resources (removed in Python 3.12 / setuptools v70+)
+# Downgrading setuptools breaks PyTorch's native C++ extensions. 
+# This mock bypasses the error safely without breaking anything else.
+mock_pkg = types.ModuleType('pkg_resources')
+mock_pkg.require = lambda *args, **kwargs: []
+mock_pkg.get_distribution = lambda *args, **kwargs: type('obj', (object,), {'version': '1.0'})()
+sys.modules['pkg_resources'] = mock_pkg
+# ----------------------------------------
 
 import math
 import shutil
 import inspect
 import zipfile
+import subprocess
 import traceback
-import importlib
 import urllib.request
 
 import numpy as np
@@ -31,9 +32,8 @@ import gradio as gr
 try:
     import spaces
     gpu = spaces.GPU(duration=60)
-except Exception:  # running locally / on CPU hardware without the `spaces` package
-    def gpu(fn):
-        return fn
+except Exception:  
+    def gpu(fn): return fn
 
 SPACE_ROOT = os.path.dirname(os.path.abspath(__file__))
 REPO = os.path.join(SPACE_ROOT, "DESED_task")
@@ -44,21 +44,15 @@ HPARAMS_PATH = os.path.join(SPACE_ROOT, "hparams.yaml")
 
 SR = 16000
 CHUNK_SEC = 10
-MAX_SEC = 60  # longer audio is truncated
-
+MAX_SEC = 60  
 
-# --------------------------------------------------------------------------- #
 # 1. Get the recipe code
-# --------------------------------------------------------------------------- #
 if not os.path.isdir(os.path.join(RECIPE, "local")):
     print("Cloning DESED_task repository...")
     shutil.rmtree(REPO, ignore_errors=True)
-    subprocess.run(
-        ["git", "clone", "--depth", "1", "https://github.com/DCASE-REPO/DESED_task.git", REPO],
-        check=True,
-    )
+    subprocess.run(["git", "clone", "--depth", "1", "https://github.com/DCASE-REPO/DESED_task.git", REPO], check=True)
 
-for p in (RECIPE, REPO):          # `local` is in RECIPE, `desed_task` is in REPO
+for p in (RECIPE, REPO):          
     if p not in sys.path:
         sys.path.insert(0, p)
 
@@ -67,12 +61,8 @@ init_file = os.path.join(local_dir, "__init__.py")
 if not os.path.exists(init_file):
     open(init_file, "w").close()
 
-print("recipe/local contents:", sorted(os.listdir(local_dir)))
-
-# --------------------------------------------------------------------------- #
-# 2. Import the recipe (SEDTask4 lives in local/sed_trainer_pretrained.py)
-# --------------------------------------------------------------------------- #
-os.chdir(RECIPE)  # some recipe modules use relative paths
+# 2. Import the recipe 
+os.chdir(RECIPE)  
 try:
     from local.sed_trainer_pretrained import SEDTask4
     from local.classes_dict import classes_labels
@@ -83,14 +73,9 @@ except Exception:
     print("CRITICAL import error while loading the DESED_task recipe:")
     traceback.print_exc()
     raise
-print("SEDTask4.__init__ signature:", inspect.signature(SEDTask4.__init__))
-
 
-# --------------------------------------------------------------------------- #
 # 3. Helpers
-# --------------------------------------------------------------------------- #
 def call_with_supported_kwargs(fn, **kwargs):
-    """Call fn with only the kwargs it accepts; required-but-unknown params get None."""
     sig = inspect.signature(fn)
     accepts_var_kw = any(p.kind == p.VAR_KEYWORD for p in sig.parameters.values())
     out = {}
@@ -100,250 +85,150 @@ def call_with_supported_kwargs(fn, **kwargs):
         if name in kwargs:
             out[name] = kwargs[name]
         elif p.default is p.empty:
-            print(f"  note: passing None for required parameter '{name}' of {fn}")
             out[name] = None
     if accepts_var_kw:
-        for k, v in kwargs.items():
-            out.setdefault(k, v)
+        for k, v in kwargs.items(): out.setdefault(k, v)
     return fn(**out)
 
-
 def torch_load(path, **kw):
-    try:
-        return torch.load(path, map_location="cpu", weights_only=True, **kw)
-    except Exception:
-        return torch.load(path, map_location="cpu", weights_only=False, **kw)
-
+    try: return torch.load(path, map_location="cpu", weights_only=True, **kw)
+    except Exception: return torch.load(path, map_location="cpu", weights_only=False, **kw)
 
 def file_ok(path, min_mb):
     return os.path.isfile(path) and os.path.getsize(path) > min_mb * 1024 * 1024
 
-
 def ensure_beats_checkpoint():
-    """BEATs_iter3_plus_AS2M.pt (self-supervised, NOT the AudioSet-finetuned one)."""
-    if file_ok(BEATS_PATH, 50):
-        return BEATS_PATH
-
-    # Option A: a copy somewhere on the Hub (third-party; loaded with weights_only=True below)
-    try:
-        from huggingface_hub import HfApi, hf_hub_download
-        for repo_id in ("tsinghua-ee/WAVE-7B",):
-            try:
-                files = [f for f in HfApi().list_repo_files(repo_id)
-                         if f.endswith("BEATs_iter3_plus_AS2M.pt")]
-                if files:
-                    print(f"Downloading {files[0]} from {repo_id} ...")
-                    path = hf_hub_download(repo_id, files[0])
-                    if file_ok(path, 50):
-                        return path
-            except Exception as e:
-                print(f"Hub download from {repo_id} failed: {e!r}")
-    except Exception as e:
-        print("huggingface_hub not usable:", repr(e))
-
-    # Option B: the original Microsoft link (often dead)
+    if file_ok(BEATS_PATH, 50): return BEATS_PATH
     try:
         url = "https://valle.blob.core.windows.net/share/BEATs/BEATs_iter3_plus_AS2M.pt"
-        print("Trying original BEATs URL ...")
         urllib.request.urlretrieve(url, BEATS_PATH)
-        if file_ok(BEATS_PATH, 50):
-            return BEATS_PATH
+        if file_ok(BEATS_PATH, 50): return BEATS_PATH
         os.remove(BEATS_PATH)
-    except Exception as e:
-        print("Original BEATs URL failed:", repr(e))
-
-    raise FileNotFoundError(
-        "BEATs_iter3_plus_AS2M.pt not found. Download it from the BEATs README "
-        "(github.com/microsoft/unilm/tree/master/beats) and upload it to the root of this Space."
-    )
-
+    except Exception: pass
+    raise FileNotFoundError("BEATs_iter3_plus_AS2M.pt not found.")
 
 def find_config(hp):
-    """The recipe's full config dict has keys like 'data', 'feats', 'net'."""
     if isinstance(hp, dict):
-        if isinstance(hp.get("hparams"), dict) and "net" in hp["hparams"]:
-            return hp["hparams"]
-        if "net" in hp and "data" in hp and "feats" in hp:
-            return hp
+        if isinstance(hp.get("hparams"), dict) and "net" in hp["hparams"]: return hp["hparams"]
+        if "net" in hp and "data" in hp and "feats" in hp: return hp
     return None
 
-
 def load_config(ckpt):
     cfg = find_config(ckpt.get("hyper_parameters"))
-    if cfg is not None:
-        print("Config taken from the checkpoint's hyper_parameters.")
-        return cfg
-    print("hyper_parameters had no usable config; falling back to hparams.yaml")
+    if cfg is not None: return cfg
     if not os.path.exists(HPARAMS_PATH):
         zip_path = os.path.join(SPACE_ROOT, "baseline_pretrained_2024.zip")
-        urllib.request.urlretrieve(
-            "https://zenodo.org/records/11280943/files/baseline_pretrained_2024.zip?download=1",
-            zip_path,
-        )
+        urllib.request.urlretrieve("https://zenodo.org/records/11280943/files/baseline_pretrained_2024.zip?download=1", zip_path)
         with zipfile.ZipFile(zip_path) as z:
             name = next(n for n in z.namelist() if n.endswith("hparams.yaml"))
             with z.open(name) as src, open(HPARAMS_PATH, "wb") as dst:
                 shutil.copyfileobj(src, dst)
         os.remove(zip_path)
-    with open(HPARAMS_PATH) as f:
-        cfg = find_config(yaml.safe_load(f))
-    if cfg is None:
-        raise RuntimeError("Could not find the recipe config (data/feats/net) in the checkpoint or hparams.yaml")
+    with open(HPARAMS_PATH) as f: cfg = find_config(yaml.safe_load(f))
     return cfg
 
-
-# --------------------------------------------------------------------------- #
 # 4. Build the model on CPU
-# --------------------------------------------------------------------------- #
 model = None
 encoder = None
 LABELS = []
 LOAD_ERROR = None
 
 try:
-    print("Loading checkpoint:", CKPT_PATH)
     ckpt = torch_load(CKPT_PATH)
     config = load_config(ckpt)
-
     encoder = call_with_supported_kwargs(
-        ManyHotEncoder,
-        labels=list(classes_labels.keys()),
-        audio_len=config["data"]["audio_max_len"],
-        frame_len=config["feats"]["n_filters"],
-        frame_hop=config["feats"]["hop_length"],
-        net_pooling=config["data"]["net_subsample"],
-        fs=config["data"]["fs"],
+        ManyHotEncoder, labels=list(classes_labels.keys()), audio_len=config["data"]["audio_max_len"],
+        frame_len=config["feats"]["n_filters"], frame_hop=config["feats"]["hop_length"],
+        net_pooling=config["data"]["net_subsample"], fs=config["data"]["fs"],
     )
     LABELS = list(encoder.labels)
-    print(f"{len(LABELS)} classes:", LABELS)
-
     sed_student = CRNN(**config["net"])
-
     beats_ckpt = torch_load(ensure_beats_checkpoint())
     beats_cfg = BEATsConfig(beats_ckpt["cfg"])
     pretrained = BEATs(beats_cfg)
     pretrained.load_state_dict(beats_ckpt["model"])
     pretrained.eval()
-    for prm in pretrained.parameters():
-        prm.requires_grad = False
-
+    for prm in pretrained.parameters(): prm.requires_grad = False
     model = call_with_supported_kwargs(
-        SEDTask4,
-        hparams=config,
-        encoder=encoder,
-        sed_student=sed_student,
-        pretrained_model=pretrained,
-        opt=None, train_data=None, valid_data=None, test_data=None,
+        SEDTask4, hparams=config, encoder=encoder, sed_student=sed_student,
+        pretrained_model=pretrained, opt=None, train_data=None, valid_data=None, test_data=None,
         train_sampler=None, scheduler=None, fast_dev_run=False, evaluation=False,
     )
-    missing, unexpected = model.load_state_dict(ckpt["state_dict"], strict=False)
-    print(f"state_dict loaded. missing={len(missing)} unexpected={len(unexpected)}")
-    print("  missing (first 10):", missing[:10])
-    print("  unexpected (first 10):", unexpected[:10])
+    model.load_state_dict(ckpt["state_dict"], strict=False)
     model.eval()
-    print("Model loaded successfully!")
 except Exception:
     LOAD_ERROR = traceback.format_exc()
-    print("Model loading error:\n", LOAD_ERROR)
 finally:
     os.chdir(SPACE_ROOT)
 
-
-# --------------------------------------------------------------------------- #
 # 5. Inference
-# --------------------------------------------------------------------------- #
 def compute_embeddings(wav):
-    """Frame-level BEATs embeddings for a (1, samples) tensor."""
     helper = getattr(model, "get_embeddings", None)
     if callable(helper):
-        try:
-            return helper(wav)
-        except TypeError:
-            pass
+        try: return helper(wav)
+        except TypeError: pass
     pm = model.pretrained_model
     out = pm.extract_features(wav, padding_mask=torch.zeros_like(wav).bool())
     return out[0] if isinstance(out, (tuple, list)) else out
 
-
 def forward_strong(wav):
-    """Returns frame probabilities as a numpy array (frames, classes)."""
     net = model.sed_student
     feats = model.mel_spec(wav)
     emb = compute_embeddings(wav)
-
     detect = model.detect
     takes_emb = "embeddings" in inspect.signature(detect).parameters
     last_err = None
-    for e in (emb, emb.transpose(1, 2)):          # the recipe's layout is not guaranteed
+    for e in (emb, emb.transpose(1, 2)):
         try:
             res = detect(feats, net, e) if takes_emb else detect(feats, net)
             strong = res[0] if isinstance(res, (tuple, list)) else res
             break
-        except Exception as err:
-            last_err = err
-    else:
-        raise last_err
-
+        except Exception as err: last_err = err
+    else: raise last_err
     p = strong[0].detach().float().cpu().numpy()
-    if p.shape[0] == len(LABELS) and p.shape[1] != len(LABELS):
-        p = p.T                                    # -> (frames, classes)
+    if p.shape[0] == len(LABELS) and p.shape[1] != len(LABELS): p = p.T
     return p
 
-
 def runs(mask):
-    """Start/end frame indices (end exclusive) of contiguous True runs."""
     d = np.diff(np.concatenate([[0], mask.astype(np.int8), [0]]))
     return list(zip(np.where(d == 1)[0], np.where(d == -1)[0]))
 
-
 @gpu
 def process_audio(audio_file, threshold):
-    if audio_file is None:
-        return "Please upload an audio file."
-    if model is None:
-        return "Model failed to load at startup:\n\n" + (LOAD_ERROR or "unknown error")
-
+    if audio_file is None: return "Please upload an audio file."
+    if model is None: return "Model failed to load:\n\n" + (LOAD_ERROR or "unknown error")
     try:
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
-
         wav, _ = librosa.load(audio_file, sr=SR, mono=True)
         wav = wav[: MAX_SEC * SR]
         chunk = CHUNK_SEC * SR
         n_chunks = max(1, math.ceil(len(wav) / chunk))
-
         events = []
         for i in range(n_chunks):
             seg = wav[i * chunk:(i + 1) * chunk]
             real_sec = len(seg) / SR
             seg = np.pad(seg, (0, chunk - len(seg)))
             x = torch.from_numpy(seg).float().unsqueeze(0).to(device)
-            with torch.no_grad():
-                probs = forward_strong(x)                      # (frames, classes)
+            with torch.no_grad(): probs = forward_strong(x)
             frame_sec = CHUNK_SEC / probs.shape[0]
             for c, name in enumerate(LABELS):
                 for s, e in runs(probs[:, c] > threshold):
                     start = i * CHUNK_SEC + s * frame_sec
                     end = i * CHUNK_SEC + min(e * frame_sec, real_sec)
-                    if end > start:
-                        events.append((start, end, name, float(probs[s:e, c].max())))
-
-        if not events:
-            return "No sound events detected at this threshold."
+                    if end > start: events.append((start, end, name, float(probs[s:e, c].max())))
+        if not events: return "No sound events detected at this threshold."
         events.sort()
         return "\n".join(f"[{s:6.2f}s - {e:6.2f}s] {n} (max prob {p:.2f})" for s, e, n, p in events)
-
     except Exception:
         return "Error during processing:\n\n" + traceback.format_exc()
     finally:
-        if model is not None:
-            model.to("cpu")
-
+        if model is not None: model.to("cpu")
 
 demo = gr.Interface(
     fn=process_audio,
     inputs=[
-        gr.Audio(type="filepath", label="Upload audio (WAV / MP3, first 60 s are used)"),
+        gr.Audio(type="filepath", label="Upload audio (WAV / MP3, max 60s)"),
         gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Detection threshold"),
     ],
     outputs=gr.Textbox(label="Detected events", lines=20),