avi292423 commited on
Commit
8b22f56
·
verified ·
1 Parent(s): ae94ca3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +59 -75
app.py CHANGED
@@ -1,75 +1,59 @@
1
- from fastapi import FastAPI, File, UploadFile
2
- from fastapi.responses import JSONResponse
3
- from models.model_wav2vec import Wav2VecIntent
4
- from huggingface_hub import hf_hub_download
5
- import torch
6
- import soundfile as sf
7
- import gradio as gr
8
- import numpy as np
9
- import librosa
10
-
11
- app = FastAPI()
12
-
13
- # Download model from Hugging Face
14
- MODEL_PATH = hf_hub_download(repo_id="avi292423/speech-intent-recognition-project", filename="wav2vec_best_model.pt")
15
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16
-
17
- label_map = {
18
- "activate_lamp": 0, "activate_lights": 1, "activate_lights_bedroom": 2, "activate_lights_kitchen": 3,
19
- "activate_lights_washroom": 4, "activate_music": 5, "bring_juice": 6, "bring_newspaper": 7,
20
- "bring_shoes": 8, "bring_socks": 9, "change_language_Chinese": 10, "change_language_English": 11,
21
- "change_language_German": 12, "change_language_Korean": 13, "change_language_none": 14,
22
- "deactivate_lamp": 15, "deactivate_lights": 16, "deactivate_lights_bedroom": 17, "deactivate_lights_kitchen": 18,
23
- "deactivate_lights_washroom": 19, "deactivate_music": 20, "decrease_heat": 21, "decrease_heat_bedroom": 22,
24
- "decrease_heat_kitchen": 23, "decrease_heat_washroom": 24, "decrease_volume": 25, "increase_heat": 26,
25
- "increase_heat_bedroom": 27, "increase_heat_kitchen": 28, "increase_heat_washroom": 29, "increase_volume": 30
26
- }
27
- index_to_label = {v: k for k, v in label_map.items()}
28
-
29
- num_classes = 31
30
- pretrained_model = "facebook/wav2vec2-base" # Use base for less RAM, or keep large if needed
31
- model = Wav2VecIntent(num_classes=num_classes, pretrained_model=pretrained_model).to(device)
32
- state_dict = torch.load(MODEL_PATH, map_location=device)
33
- model.load_state_dict(state_dict)
34
- model.eval()
35
-
36
- @app.post("/predict")
37
- async def predict(file: UploadFile = File(...)):
38
- audio_bytes = await file.read()
39
- with open("temp.wav", "wb") as f:
40
- f.write(audio_bytes)
41
- audio, sample_rate = sf.read("temp.wav")
42
- if sample_rate != 16000:
43
- return JSONResponse({"error": "Audio must have a sample rate of 16kHz."}, status_code=400)
44
- waveform = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
45
- with torch.no_grad():
46
- output = model(waveform)
47
- predicted_class = torch.argmax(output, dim=1).item()
48
- predicted_label = index_to_label.get(predicted_class, "Unknown Class")
49
- return {"prediction": predicted_label}
50
-
51
- def predict_intent(audio):
52
- if audio is None:
53
- return "No audio provided."
54
- sr, y = audio
55
- if sr != 16000:
56
- # Resample to 16kHz
57
- y = librosa.resample(y.astype(float), orig_sr=sr, target_sr=16000)
58
- sr = 16000
59
- waveform = torch.tensor(y, dtype=torch.float32).unsqueeze(0).to(device)
60
- with torch.no_grad():
61
- output = model(waveform)
62
- predicted_class = torch.argmax(output, dim=1).item()
63
- predicted_label = index_to_label.get(predicted_class, "Unknown Class")
64
- return predicted_label
65
-
66
- demo = gr.Interface(
67
- fn=predict_intent,
68
- inputs=gr.Audio(source="microphone", type="numpy", label="Record or Upload Audio (16kHz WAV)"),
69
- outputs=gr.Textbox(label="Predicted Intent"),
70
- title="Speech Intent Recognition",
71
- description="Record or upload a 16kHz WAV audio file to predict the intent."
72
- )
73
-
74
- if __name__ == "__main__":
75
- demo.launch()
 
1
+ from fastapi import FastAPI, File, UploadFile
2
+ from fastapi.responses import JSONResponse
3
+ from fastapi.middleware.cors import CORSMiddleware
4
+ from models.model_wav2vec import Wav2VecIntent
5
+ from huggingface_hub import hf_hub_download
6
+ import torch
7
+ import soundfile as sf
8
+ import numpy as np
9
+ import librosa
10
+
11
+ app = FastAPI()
12
+
13
+ # Enable CORS for all origins (so your frontend can call the API)
14
+ app.add_middleware(
15
+ CORSMiddleware,
16
+ allow_origins=["*"],
17
+ allow_credentials=True,
18
+ allow_methods=["*"],
19
+ allow_headers=["*"],
20
+ )
21
+
22
+ # Download model from Hugging Face
23
+ MODEL_PATH = hf_hub_download(repo_id="avi292423/speech-intent-recognition-project", filename="wav2vec_best_model.pt")
24
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
25
+
26
+ label_map = {
27
+ "activate_lamp": 0, "activate_lights": 1, "activate_lights_bedroom": 2, "activate_lights_kitchen": 3,
28
+ "activate_lights_washroom": 4, "activate_music": 5, "bring_juice": 6, "bring_newspaper": 7,
29
+ "bring_shoes": 8, "bring_socks": 9, "change_language_Chinese": 10, "change_language_English": 11,
30
+ "change_language_German": 12, "change_language_Korean": 13, "change_language_none": 14,
31
+ "deactivate_lamp": 15, "deactivate_lights": 16, "deactivate_lights_bedroom": 17, "deactivate_lights_kitchen": 18,
32
+ "deactivate_lights_washroom": 19, "deactivate_music": 20, "decrease_heat": 21, "decrease_heat_bedroom": 22,
33
+ "decrease_heat_kitchen": 23, "decrease_heat_washroom": 24, "decrease_volume": 25, "increase_heat": 26,
34
+ "increase_heat_bedroom": 27, "increase_heat_kitchen": 28, "increase_heat_washroom": 29, "increase_volume": 30
35
+ }
36
+ index_to_label = {v: k for k, v in label_map.items()}
37
+
38
+ num_classes = 31
39
+ pretrained_model = "facebook/wav2vec2-large" # Use large model
40
+ model = Wav2VecIntent(num_classes=num_classes, pretrained_model=pretrained_model).to(device)
41
+ state_dict = torch.load(MODEL_PATH, map_location=device)
42
+ model.load_state_dict(state_dict)
43
+ model.eval()
44
+
45
+ @app.post("/predict")
46
+ async def predict(file: UploadFile = File(...)):
47
+ audio_bytes = await file.read()
48
+ with open("temp.wav", "wb") as f:
49
+ f.write(audio_bytes)
50
+ audio, sample_rate = sf.read("temp.wav")
51
+ if sample_rate != 16000:
52
+ # Resample to 16kHz
53
+ audio = librosa.resample(audio.astype(float), orig_sr=sample_rate, target_sr=16000)
54
+ waveform = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
55
+ with torch.no_grad():
56
+ output = model(waveform)
57
+ predicted_class = torch.argmax(output, dim=1).item()
58
+ predicted_label = index_to_label.get(predicted_class, "Unknown Class")
59
+ return {"prediction": predicted_label}