sawadogosalif commited on
Commit
4fdf0a0
·
verified ·
1 Parent(s): ae31542

Create preprocess.py

Browse files
Files changed (1) hide show
  1. preprocess.py +87 -0
preprocess.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from typing import Tuple, Dict
3
+ from joblib import Parallel, delayed
4
+ from resemble_enhance.enhancer.inference import denoise, enhance
5
+
6
+
7
+ def denoise_audio(audio: torch.Tensor, sr: int, device: str) -> Tuple[torch.Tensor, int]:
8
+ """
9
+ Réduction de bruit sur un tenseur audio.
10
+ """
11
+ try:
12
+ return denoise(audio, sr, device)
13
+ except Exception as e:
14
+ print(f"[denoise_audio] Erreur: {e}")
15
+ return audio, sr
16
+
17
+
18
+ def enhance_audio(
19
+ audio: torch.Tensor,
20
+ sr: int,
21
+ device: str,
22
+ solver: str = "midpoint",
23
+ nfe: int = 128,
24
+ lambd: float = 0.1,
25
+ tau: float = 0.01
26
+ ) -> Tuple[torch.Tensor, int]:
27
+ """
28
+ Amélioration audio (résolution d'ODE).
29
+ """
30
+ try:
31
+ return enhance(
32
+ audio,
33
+ sr,
34
+ device,
35
+ solver=solver.lower(),
36
+ nfe=nfe,
37
+ lambd=lambd,
38
+ tau=tau
39
+ )
40
+ except Exception as e:
41
+ print(f"[enhance_audio] Erreur: {e}")
42
+ return audio, sr
43
+
44
+
45
+ def process_audio(
46
+ audio: torch.Tensor,
47
+ sr: int,
48
+ device: str = "cpu",
49
+ solver: str = "midpoint",
50
+ nfe: int = 128,
51
+ tau: float = 0.01,
52
+ denoise_before: bool = False,
53
+ ) -> Dict[str, Tuple[torch.Tensor, int]]:
54
+ """
55
+ Pipeline complet de traitement audio.
56
+
57
+ Args:
58
+ audio: Tensor d'entrée de forme (1, T).
59
+ sr: Fréquence d'échantillonnage.
60
+ device: 'cpu' ou 'cuda'.
61
+ solver: ODE solver.
62
+ nfe: nombre d'évaluations de la fonction.
63
+ tau: paramètre tau.
64
+ denoise_before: si True, lambd=0.9, sinon lambd=0.1.
65
+ parallel: si True, exécute denoise et enhance en parallèle (utile si indépendants).
66
+
67
+ Returns:
68
+ dict avec clés 'denoised' et 'enhanced', valeurs (tensor, sr).
69
+ """
70
+ # Choix de lambda
71
+ lambd = 0.9 if denoise_before else 0.1
72
+ # Exécution séquentielle (débruitage puis amélioration)
73
+ denoised = denoise_audio(audio, sr, device)
74
+ # Applique enhance sur le résultat du débruitage
75
+ enhanced = enhance_audio(
76
+ denoised[0],
77
+ denoised[1],
78
+ device,
79
+ solver,
80
+ nfe,
81
+ lambd,
82
+ tau
83
+ )
84
+ return {
85
+ "denoised": denoised,
86
+ "enhanced": enhanced
87
+ }