Update app.py

app.py

@@ -383,6 +383,7 @@ def build_encoder(config):
 # 5. Load the model on CPU
 # --------------------------------------------------
 model = None
+pretrained_model = None  # Kept separate so its weights are not part of the SED checkpoint.
 encoder = None
 LABELS = []
 LOAD_ERROR = None
@@ -398,6 +399,12 @@ try:
 
     config = load_config(ckpt)
 
+    if config.get("pretrained", {}).get("e2e", False):
+        raise ValueError(
+            "This app currently targets the official DCASE 2024 baseline "
+            "with pretrained.e2e=False and precomputed BEATs embeddings."
+        )
+
     # Build the correctly ordered 27-class encoder.
     encoder = build_encoder(config)
     LABELS = list(encoder.labels)
@@ -405,29 +412,61 @@ try:
     # Build the CRNN student model.
     sed_student = CRNN(**config["net"])
 
-    # Load the pretrained BEATs model.
+    # Load only the trained student weights. The official DCASE recipe saves
+    # this submodule separately in checkpoint["sed_student"]. Loading the
+    # entire Lightning state_dict with strict=False can silently leave the
+    # detector partly uninitialized, especially across recipe revisions.
+    student_state = ckpt.get("sed_student")
+
+    if not isinstance(student_state, dict):
+        full_state = ckpt.get("state_dict", {})
+        if isinstance(full_state, dict):
+            prefixes = ("sed_student.", "module.sed_student.")
+            for prefix in prefixes:
+                extracted = {
+                    key[len(prefix):]: value
+                    for key, value in full_state.items()
+                    if key.startswith(prefix)
+                }
+                if extracted:
+                    student_state = extracted
+                    break
+
+    if not isinstance(student_state, dict) or not student_state:
+        raise ValueError(
+            "Could not find trained student weights in checkpoint. "
+            "Expected checkpoint['sed_student'] or keys prefixed with 'sed_student.'."
+        )
+
+    # Strict loading is intentional: fail rather than serve random/partial weights.
+    sed_student.load_state_dict(student_state, strict=True)
+    print(
+        f"Loaded student checkpoint strictly ({len(student_state)} tensors)."
+    )
+
+    # Load BEATs separately for online waveform embedding extraction.
+    # The official DCASE 2024 baseline uses precomputed BEATs embeddings
+    # (pretrained.e2e=False), so BEATs must not be attached to the SEDTask4
+    # module or included in its checkpoint state_dict.
     beats_path = ensure_beats_checkpoint()
     beats_ckpt = torch_load(beats_path)
 
     beats_cfg = BEATsConfig(beats_ckpt["cfg"])
-    pretrained = BEATs(beats_cfg)
-
-    pretrained.load_state_dict(
-        beats_ckpt["model"]
-    )
-
-    pretrained.eval()
+    pretrained_model = BEATs(beats_cfg)
+    pretrained_model.load_state_dict(beats_ckpt["model"], strict=True)
+    pretrained_model.eval()
 
-    for parameter in pretrained.parameters():
+    for parameter in pretrained_model.parameters():
         parameter.requires_grad = False
 
-    # Construct the DCASE sound event detection model.
+    # Construct the DCASE detector. It deep-copies the already loaded student
+    # into sed_teacher; inference below uses sed_student only.
     model = call_with_supported_kwargs(
         SEDTask4,
         hparams=config,
         encoder=encoder,
         sed_student=sed_student,
-        pretrained_model=pretrained,
+        pretrained_model=None,
         opt=None,
         train_data=None,
         valid_data=None,
@@ -438,33 +477,6 @@ try:
         evaluation=False,
     )
 
-    # Ensure raw-audio inference can access BEATs even when
-    # the checkpoint configuration uses precomputed embeddings.
-    model.pretrained_model = pretrained
-
-    # Load the trained detector weights.
-    state_dict = ckpt.get("state_dict")
-
-    if not isinstance(state_dict, dict):
-        raise ValueError(
-            "The checkpoint does not contain a valid state_dict."
-        )
-
-    incompat = model.load_state_dict(
-        state_dict,
-        strict=False,
-    )
-
-    print("Checkpoint loaded.")
-    print(
-        "Missing state-dict keys:",
-        len(incompat.missing_keys),
-    )
-    print(
-        "Unexpected state-dict keys:",
-        len(incompat.unexpected_keys),
-    )
-
     model.eval()
 
     # Quick consistency check.
@@ -493,45 +505,18 @@ finally:
 # 6. Inference helpers
 # --------------------------------------------------
 def compute_embeddings(wav):
-    """
-    Extract frame-level BEATs embeddings.
+    """Extract frame-level embeddings in the format used by DCASE's BEATsModel.
 
-    wav shape: [batch_size, samples]
-    Expected embedding output: [batch_size, frames, features]
+    Input waveform: [batch, samples].
+    Return embeddings: [batch, embedding_channels, frames].
     """
-    helper = getattr(model, "get_embeddings", None)
-
-    if callable(helper):
-        try:
-            return helper(wav)
-        except TypeError:
-            pass
-
-    pretrained_model = getattr(
-        model,
-        "pretrained_model",
-        None,
-    )
-
     if pretrained_model is None:
-        raise RuntimeError(
-            "The model has no pretrained BEATs encoder."
-        )
+        raise RuntimeError("The BEATs encoder has not been loaded.")
 
-    padding_mask = torch.zeros_like(
-        wav,
-        dtype=torch.bool,
-    )
-
-    result = pretrained_model.extract_features(
-        wav,
-        padding_mask=padding_mask,
-    )
-
-    if isinstance(result, (tuple, list)):
-        return result[0]
-
-    return result
+    # Match local/beats/BEATs.py -> BEATsModel.forward():
+    # BEATs returns [batch, frames, 768]; the detector expects [batch, 768, frames].
+    features, _ = pretrained_model.extract_features(wav)
+    return features.transpose(1, 2).float()
 
 
 def forward_strong(wav):
@@ -558,32 +543,14 @@ def forward_strong(wav):
         takes_embeddings = True
 
     if takes_embeddings:
-        last_error = None
-
-        # Support recipe variants that expect different embedding axes.
-        for candidate in (
-            embeddings,
-            embeddings.transpose(1, 2),
-        ):
-            try:
-                result = detect(
-                    feats,
-                    net,
-                    embeddings=candidate,
-                )
-
-                break
-
-            except Exception as exc:
-                last_error = exc
-
-        else:
-            raise RuntimeError(
-                "Detector failed with both embedding layouts."
-            ) from last_error
-
+        # compute_embeddings already returns [batch, 768, frames], which is
+        # the orientation of the saved DCASE frame_embeddings arrays.
+        result = detect(feats, net, embeddings=embeddings)
     else:
-        result = detect(feats, net)
+        raise RuntimeError(
+            "This DCASE detect() method does not accept embeddings; "
+            "the current checkpoint requires BEATs embeddings."
+        )
 
     # Some recipe versions return (strong_predictions, other_output).
     strong = (
@@ -662,6 +629,8 @@ def process_audio(audio_file, threshold):
         )
 
         model.to(device)
+        if pretrained_model is not None:
+            pretrained_model.to(device)
 
         # Load mono audio at 16 kHz.
         wav, _ = librosa.load(
@@ -787,6 +756,11 @@ def process_audio(audio_file, threshold):
                 model.to("cpu")
             except Exception:
                 pass
+        if pretrained_model is not None:
+            try:
+                pretrained_model.to("cpu")
+            except Exception:
+                pass
 
 
 # --------------------------------------------------