File size: 6,811 Bytes
fb4f813 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 |
# Новое содержимое для app/api.py
api_content = '''import fastapi
import shutil
import os
import zipfile
import io
import uvicorn
import glob
from typing import List
import torch
import numpy as np
import soundfile as sf
class ModelAPI:
def __init__(self, host, port):
self.host = host
self.port = port
self.base_path = os.path.join(os.path.expanduser("~"), ".modelapi")
self.noisy_audio_path = os.path.join(self.base_path, "noisy_audio")
self.enhanced_audio_path = os.path.join(self.base_path, "enhanced_audio")
# Model parameters
self.model = None
self.device = "cpu" # Force CPU since no GPU
# Create directories if they do not exist
for audio_path in [self.noisy_audio_path, self.enhanced_audio_path]:
if not os.path.exists(audio_path):
os.makedirs(audio_path)
# Clean directories
for filename in os.listdir(audio_path):
file_path = os.path.join(audio_path, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
elif os.path.isdir(file_path):
shutil.rmtree(file_path)
except Exception as e:
raise e
self.app = fastapi.FastAPI()
self._setup_routes()
def _prepare(self):
"""Initialize the speech enhancement model"""
try:
# Import here to avoid loading during container build
from speechbrain.pretrained import SepformerSeparation
print(f"Loading SpeechBrain model on {self.device}...")
self.model = SepformerSeparation.from_hparams(
source="loveisgone/sepformer-wham-enhancement-no1",
savedir="pretrained_models/sepformer-wham",
run_opts={"device": self.device}
)
print("Model loaded successfully!")
except Exception as e:
print(f"Error loading SpeechBrain: {e}")
# Fallback to simple denoising
print("Using simple denoising method")
self.model = "simple"
def _enhance(self):
"""Enhance audio files"""
noisy_files = sorted(glob.glob(os.path.join(self.noisy_audio_path, "*.wav")))
for noisy_file in noisy_files:
try:
if self.model == "simple":
# Simple noise reduction
audio, sr = sf.read(noisy_file)
# Ensure mono
if len(audio.shape) > 1:
audio = np.mean(audio, axis=1)
# Simple high-pass filter to remove low-frequency noise
from scipy import signal
b, a = signal.butter(4, 100/(sr/2), "high")
enhanced = signal.filtfilt(b, a, audio)
# Normalize
max_val = np.max(np.abs(enhanced))
if max_val > 0:
enhanced = enhanced / max_val * 0.95
else:
# Use SpeechBrain model
enhanced = self.model.separate_file(path=noisy_file)
enhanced = enhanced[:, 0].cpu().numpy()
# Get sample rate
_, sr = sf.read(noisy_file)
# Save enhanced audio
output_path = os.path.join(
self.enhanced_audio_path,
os.path.basename(noisy_file)
)
sf.write(output_path, enhanced, sr)
print(f"Enhanced: {os.path.basename(noisy_file)}")
except Exception as e:
print(f"Error processing {noisy_file}: {e}")
raise e
def _setup_routes(self):
"""Setup API routes"""
self.app.get("/status/")(self.get_status)
self.app.post("/prepare/")(self.prepare)
self.app.post("/upload-audio/")(self.upload_audio)
self.app.post("/enhance/")(self.enhance_audio)
self.app.get("/download-enhanced/")(self.download_enhanced)
async def get_status(self):
try:
return {"container_running": True}
except:
raise fastapi.HTTPException(status_code=500, detail="An error occurred while fetching API status.")
async def prepare(self):
try:
self._prepare()
return {"preparations": True}
except Exception as e:
return {"preparations": False, "error": str(e)}
async def upload_audio(self, files: List[fastapi.UploadFile] = fastapi.File(...)):
uploaded_files = []
for file in files:
try:
file_path = os.path.join(self.noisy_audio_path, file.filename)
content = await file.read()
with open(file_path, "wb") as f:
f.write(content)
uploaded_files.append(file.filename)
except Exception as e:
raise fastapi.HTTPException(status_code=500, detail=f"An error occurred while uploading: {e}")
return {"uploaded_files": uploaded_files, "status": True}
async def enhance_audio(self):
try:
self._enhance()
return {"status": True}
except Exception as e:
raise fastapi.HTTPException(status_code=500, detail=f"An error occurred while enhancing: {e}")
async def download_enhanced(self):
try:
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w") as zip_file:
for wav_file in glob.glob(os.path.join(self.enhanced_audio_path, "*.wav")):
zip_file.write(wav_file, arcname=os.path.basename(wav_file))
zip_buffer.seek(0)
return fastapi.responses.StreamingResponse(
iter([zip_buffer.getvalue()]),
media_type="application/zip",
headers={"Content-Disposition": "attachment; filename=enhanced_audio_files.zip"}
)
except Exception as e:
raise fastapi.HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
def run(self):
uvicorn.run(self.app, host=self.host, port=self.port)
'''
with open('app/api.py', 'w') as f:
f.write(api_content)
print("API updated successfully!")
|