Update app.py

app.py

@@ -1,76 +1,93 @@
-import gradio as gr, pandas as pd, zipfile, tempfile, shutil, pathlib, torch
-from utmosv2_batch_predict import compute_spec, MAX_LEN         # reuse the function
-from utmosv2.utils          import get_model
+import gradio as gr
+import torch, torchaudio, importlib, pandas as pd, tempfile, zipfile, pathlib, shutil, numpy as np
 from types import SimpleNamespace
-import importlib, torchaudio, numpy as np, torch.nn as nn
+from utmosv2.utils import get_model
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MAX_LEN = 160_000            # 10 s @16 kHz
 
-# --- load UTMOSv2 once ---
+# ---- mudel laetakse korra kogu Space’i eluea jooksul ----
 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 = 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 = "YOUR_FINE_TUNED_WEIGHT.ckpt"           # put the file in the repo
-model      = get_model(cfg, DEVICE).eval()
-specs_cfg  = cfg.dataset.specs
+cfg.weight = "YOUR_FINE_TUNED_WEIGHT.ckpt"   # lisad checks-pointi repo „Files“ vaatesse
+model = get_model(cfg, DEVICE).eval()
+specs_cfg = cfg.dataset.specs
 
-def run_space(csv_file, wav_zip, num_domains):
-    """
-    Inputs:
-        csv_file – csv with 'audio' and optional 'method'
-        wav_zip  – zip that contains all referenced .wav files
-    Output:
-        DataFrame shown + downloadable CSV
-    """
-    # ----- prepare wav directory -----
-    tempdir = tempfile.mkdtemp()
-    with zipfile.ZipFile(wav_zip.name) as zf:
-        zf.extractall(tempdir)
+def compute_spec(wav: torch.Tensor):
+    views = []
+    for s in specs_cfg:
+        mel = torchaudio.transforms.MelSpectrogram(
+            sample_rate=16000, 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)
+    return torch.stack(views)
 
-    df   = pd.read_csv(csv_file.name)
-    pred = []
-    for relpath in df["audio"]:
-        path = pathlib.Path(tempdir) / relpath
-        wav, sr = torchaudio.load(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 nn.functional.pad(wav,(0,MAX_LEN-wav.numel()))).to(DEVICE)
+def single_predict(audio_path, domain, quick):
+    # identne loogika sarulab-speech Space’iga :contentReference[oaicite:0]{index=0}
+    wav, sr = torchaudio.load(audio_path)
+    if sr != 16000:
+        wav = torchaudio.transforms.Resample(sr, 16000)(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)
+    spec = compute_spec(wav)
+    preds = []
+    # kui sul on fine-tuningus kasutatud nt 8 domääni, muuda siit
+    NUM_DOMAINS = 8
+    for dom in range(NUM_DOMAINS):
+        d_oh = torch.nn.functional.one_hot(torch.tensor(dom, device=DEVICE),
+                                           num_classes=model.num_dataset).float()[None]
+        with torch.no_grad():
+            p = model(wav[None], spec[None], d_oh).item()
+        preds.append(p)
+        if quick: break
+    return float(np.mean(preds))
 
-        spec  = compute_spec(wav, specs_cfg, DEVICE)
-        dom_p = []
-        for dom in range(int(num_domains)):
-            dom_oh = torch.nn.functional.one_hot(
-                        torch.tensor(dom,device=DEVICE),
-                        num_classes=model.num_dataset).float()[None]
-            with torch.no_grad():
-                dom_p.append(model(wav[None], spec[None], dom_oh).item())
-        pred.append(float(np.mean(dom_p)))
+def batch_predict(csv_file, wav_zip, num_domains):
+    tdir = tempfile.mkdtemp()
+    with zipfile.ZipFile(wav_zip.name) as zf:
+        zf.extractall(tdir)
+    df = pd.read_csv(csv_file.name)
+    outs = []
+    for rel in df["audio"]:
+        path = pathlib.Path(tdir) / rel
+        outs.append(single_predict(str(path), "dummy", quick=True))   # domeeni-väärtus ei loe siin
+    df["pred_mos"] = outs
+    out_file = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
+    df.to_csv(out_file, index=False)
+    shutil.rmtree(tdir)
+    return df, out_file
 
-    shutil.rmtree(tempdir)
-    df["pred_mos"] = pred
-    out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
-    df.to_csv(out_path, index=False)
-    return df, out_path                      # gr.File returns a link automatically
+with gr.Blocks(title="UTMOS-v2 MOS-hinnang") as demo:
+    gr.Markdown(
+        """
+        # UTMOS-v2  
+        Laadi üksik `.wav` või kogu partii ning saa ennustatud MOS-id.  
+        Mudel laetakse GPU-le ühe korra, seega järgmised päringud on kiiremad.
+        """
+    )
+    with gr.Tab("Üksik klipp"):
+        audio = gr.Audio(type="filepath", label="Helifail (16 kHz WAV)")
+        domain = gr.Dropdown(["default"], value="default",
+                             label="Domään (valikuline, kui ise muutsid koodi)")
+        quick = gr.Checkbox(value=True, label="Kiire (1 iteratsioon/fold)")
+        out_mos = gr.Number(label="Ennustatud MOS")
+        gr.Button("Hinda").click(single_predict, [audio, domain, quick], out_mos)
 
-demo = gr.Interface(
-    run_space,
-    inputs=[
-        gr.File(label="CSV (audio, method, MOS)", file_types=[".csv"]),
-        gr.File(label="ZIP with .wav files",       file_types=[".zip"]),
-        gr.Number(label="Training domains", value=8, precision=0)
-    ],
-    outputs=[
-        gr.Dataframe(label="Results"),
-        gr.File(label="Download predictions CSV")
-    ],
-    title="UTMOS-v2 MOS Estimator",
-    description="Upload the ground-truth CSV and a ZIP containing all WAVs. "
-                "The Space appends predicted MOS scores."
-)
+    with gr.Tab("Partii (CSV + ZIP)"):
+        csv_in  = gr.File(file_types=[".csv"], label="CSV (audio[, MOS, method])")
+        zip_in  = gr.File(file_types=[".zip"], label="ZIP kõikide WAV-idega")
+        n_dom   = gr.Number(value=8, precision=0, label="Treeningu domäänide arv")
+        df_out  = gr.Dataframe(label="Tulemused")
+        file_dl = gr.File(label="Lae CSV ennustustega")
+        gr.Button("Start").click(batch_predict,
+                                 [csv_in, zip_in, n_dom],
+                                 [df_out, file_dl])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.queue(max_size=10).launch()