Update app.py

app.py

@@ -1,4 +1,4 @@
-# ---------------- monkey-patch so CUDA checkpoints load on CPU --------------
+# ---------------- monkey-patch: CUDA-checkpoint → CPU -----------------------
 import torch
 _original_torch_load = torch.load
 def _cpu_load(*args, **kwargs):
@@ -18,94 +18,127 @@ torch.set_grad_enabled(False)
 if DEVICE.type == "cpu":
     torch.set_num_threads(min(2, os.cpu_count() or 2))
 
-MAX_LEN      = 160_000       # 10 s @16 kHz
-NUM_DOMAINS  = 3             # ← ennustame alati 3 domeeni põhjal
+MAX_LEN     = 160_000   # 10 s @16 kHz
+NUM_DOMAINS = 3         # ennustame alati 3 domääni põhjal
 # ---------------------------------------------------------------------------
 
-# ---------- load model once -------------------------------------------------
+# ---------- laadime mudeli korra -------------------------------------------
 cfg_mod = importlib.import_module("utmosv2.config.fusion_stage3")
 cfg = SimpleNamespace(**{k: getattr(cfg_mod, k) for k in dir(cfg_mod)
                          if not k.startswith("__")})
 cfg.phase, cfg.data_config, cfg.print_config = "test", None, False
 cfg.weight = "utmosv2_estonian.pth"
-model      = get_model(cfg, DEVICE).to(DEVICE).eval()
-specs_cfg  = cfg.dataset.specs
+model     = get_model(cfg, DEVICE).to(DEVICE).eval()
+specs_cfg = cfg.dataset.specs
 # ---------------------------------------------------------------------------
 
 def compute_spec(wav: torch.Tensor) -> torch.Tensor:
+    """Tagastab [V,3,512,512] multi-view mel-spectrogrammi."""
     views = []
     for s in specs_cfg:
         mel = torchaudio.transforms.MelSpectrogram(
-                sample_rate=16_000, n_fft=s.n_fft,
-                hop_length=s.hop_length, win_length=s.win_length,
-                n_mels=s.n_mels).to(DEVICE)
+            sample_rate=16_000, n_fft=s.n_fft,
+            hop_length=s.hop_length, win_length=s.win_length,
+            n_mels=s.n_mels).to(DEVICE)
         db = torchaudio.transforms.AmplitudeToDB()(mel(wav[None]))[0]
         if db.shape != (512, 512):
             db = torch.nn.functional.interpolate(db[None, None],
                                                  size=(512, 512),
                                                  mode="bilinear",
                                                  align_corners=False)[0, 0]
-        views.extend([db.repeat(3, 1, 1)] * 2)      # kaks vaadet
-    return torch.stack(views)                       # [V,3,512,512]
+        views.extend([db.repeat(3, 1, 1)] * 2)
+    return torch.stack(views)
 
 def single_predict(audio_path: str) -> float:
-    """MOS ühe faili kohta – keskmista 3 domeeni."""
+    """MOS ühe faili kohta – keskmistab 3 domääni."""
     wav, sr = torchaudio.load(audio_path)
     if sr != 16_000:
         wav = torchaudio.transforms.Resample(sr, 16_000)(wav)[0]
     else:
         wav = wav[0]
 
-    wav  = (wav[:MAX_LEN] if wav.numel() > MAX_LEN
-            else torch.nn.functional.pad(wav, (0, MAX_LEN - wav.numel()))
-           ).to(DEVICE)
+    wav = (wav[:MAX_LEN] if wav.numel() > MAX_LEN
+           else torch.nn.functional.pad(wav, (0, MAX_LEN - wav.numel()))
+          ).to(DEVICE)
     spec = compute_spec(wav)
 
     preds = []
     for dom in range(NUM_DOMAINS):
         dom_oh = torch.nn.functional.one_hot(
-                     torch.tensor(dom, device=DEVICE),
-                     num_classes=model.num_dataset
-                 ).float()[None]
+            torch.tensor(dom, device=DEVICE),
+            num_classes=model.num_dataset
+        ).float()[None]
         preds.append(model(wav[None], spec[None], dom_oh).item())
     return float(np.mean(preds))
 
+# ---------- abifunktsioon faili otsimiseks ---------------------------------
+def locate_audio(root: pathlib.Path, rel_path: str) -> pathlib.Path:
+    """
+    Leia helifail ajutisest kataloogist.
+    1) root / rel_path
+    2) kui ei leidu, otsi failinime järgi kogu puust.
+    """
+    rel_path = rel_path.strip().lstrip("/\\")
+    direct   = root / rel_path
+    if direct.is_file():
+        return direct
+
+    matches = list(root.rglob(pathlib.Path(rel_path).name))
+    if matches:
+        return matches[0]
+
+    raise FileNotFoundError(f"'{rel_path}' ei leitud ZIP-ist – "
+                            f"kontrolli CSV ja ZIP teede vastavust.")
+
+# ---------- partii-töötlus --------------------------------------------------
 def batch_predict(csv_file, wav_zip):
-    """Partii ennustused: loe CSV, lae ZIPist WAV-id, lisa 'pred_mos' veerg."""
+    """Loeb CSV & ZIP, lisab 'pred_mos', annab CSV-i ja tabeli tagasi."""
     tdir = tempfile.mkdtemp()
     with zipfile.ZipFile(wav_zip.name) as zf:
         zf.extractall(tdir)
 
-    df, outs = pd.read_csv(csv_file.name), []
-    for rel in df["faili_uus_nimi"]:
-        outs.append(single_predict(str(pathlib.Path(tdir) / rel)))
+    df   = pd.read_csv(csv_file.name)
+    col  = "faili_uus_nimi" if "faili_uus_nimi" in df.columns else "audio"
+    outs, errors = [], []
+
+    for rel in df[col]:
+        try:
+            full = locate_audio(pathlib.Path(tdir), str(rel))
+            outs.append(single_predict(str(full)))
+        except Exception as e:
+            outs.append(np.nan)
+            errors.append(f"{rel}: {e}")
 
     df["pred_mos"] = outs
     out_file = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
     df.to_csv(out_file, index=False)
     shutil.rmtree(tdir)
+
+    if errors:
+        gr.Warning("⚠️ Mõned failid jäid leidmata või tekkis viga:\n" +
+                   "\n".join(errors[:15]))
+
     return df, out_file
 # ---------------------------------------------------------------------------
 
+# --------------------------- Gradio UI -------------------------------------
 with gr.Blocks(title="UTMOS-v2 MOS-hinnang (3 domeeni)") as demo:
     gr.Markdown(
         """
         # UTMOS-v2 (eesti kõne)  
-        Ennustab objektiivse MOS-i **kolme treeningu‐domääni** keskmisena.
+        Ennustab objektiivse MOS-i **kolme treeningu-domääni** keskmisena.
         """
     )
 
-    # ---- üksik klipp -------------------------------------------------------
+    # Üksik klipp -----------------------------------------------------------
     with gr.Tab("Üksik klipp"):
-        audio   = gr.Audio(type="filepath", label="WAV (16 kHz)")
+        audio   = gr.Audio(type="filepath", label="WAV (16 kHz või muu)")
         out_mos = gr.Number(label="Ennustatud MOS")
-        gr.Button("Hinda").click(fn=single_predict,
-                                 inputs=audio,
-                                 outputs=out_mos)
+        gr.Button("Hinda").click(single_predict, inputs=audio, outputs=out_mos)
 
-    # ---- partii ------------------------------------------------------------
+    # Partii ---------------------------------------------------------------
     with gr.Tab("Partii (CSV + ZIP)"):
-        csv_in  = gr.File(file_types=[".csv"], label="CSV (audio[, MOS, method])")
+        csv_in  = gr.File(file_types=[".csv"], label="CSV (audio|faili_uus_nimi)")
         zip_in  = gr.File(file_types=[".zip"], label="ZIP kõigi WAV-idega")
         df_out  = gr.Dataframe(label="Tulemused")
         file_dl = gr.File(label="Lae CSV ennustustega")