Upload filtravimas.py

filtravimas.py

@@ -0,0 +1,86 @@
+# filtravimas.py
+
+import os
+import torch
+import torchaudio
+import torch.nn as nn
+import torchaudio.transforms as T
+from asteroid.models import DCCRNet
+
+# Laikinas katalogas išfiltruotiems failams
+TEMP_DIR = "temp_filtered"
+OUTPUT_PATH = os.path.join(TEMP_DIR, "ivestis.wav")
+
+# Užtikriname, kad aplankas egzistuoja
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+# Paprastas Wave-U-Net modelis
+class WaveUNet(nn.Module):
+    def __init__(self):
+        super(WaveUNet, self).__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv1d(1, 16, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv1d(16, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv1d(32, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose1d(64, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose1d(32, 16, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose1d(16, 1, kernel_size=3, stride=1, padding=1)
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+# Wave-U-Net filtravimas
+def filtruoti_su_waveunet(input_path, output_path):
+    print("🔧 Wave-U-Net filtravimas...")
+    model = WaveUNet()
+    model.eval()
+    mixture, sr = torchaudio.load(input_path)
+
+    if sr != 16000:
+        print("🔁 Resample į 16kHz...")
+        mixture = T.Resample(orig_freq=sr, new_freq=16000)(mixture)
+
+    mixture = mixture.unsqueeze(0)
+    with torch.no_grad():
+        output = model(mixture)
+
+    output = output.squeeze(0)
+    torchaudio.save(output_path, output, 16000)
+    print(f"✅ Išsaugotas Wave-U-Net rezultatas: {output_path}")
+
+# Denoiser (DCCRNet)
+def filtruoti_su_denoiser(input_path, output_path):
+    print("🔧 Denoiser (DCCRNet)...")
+    model = DCCRNet.from_pretrained("JorisCos/DCCRNet_Libri1Mix_enhsingle_16k")
+    mixture, sr = torchaudio.load(input_path)
+
+    if sr != 16000:
+        print("🔁 Resample į 16kHz...")
+        mixture = T.Resample(orig_freq=sr, new_freq=16000)(mixture)
+
+    with torch.no_grad():
+        est_source = model.separate(mixture)
+
+    torchaudio.save(output_path, est_source[0], 16000)
+    print(f"✅ Išsaugotas Denoiser rezultatas: {output_path}")
+
+# Pasirinkimo funkcija
+def filtruoti_audio(input_path: str, metodas: str) -> str:
+    if metodas == "Denoiser":
+        filtruoti_su_denoiser(input_path, OUTPUT_PATH)
+    elif metodas == "Wave-U-Net":
+        filtruoti_su_waveunet(input_path, OUTPUT_PATH)
+    else:
+        raise ValueError("Nepalaikomas filtravimo metodas")
+
+    return OUTPUT_PATH