quan030106/vietnamese-dialect-api

1import torch
2import torch.nn as nn
3import torchaudio
4import librosa
5import numpy as np
6import os
7from fastapi import FastAPI, File, UploadFile
8from fastapi.middleware.cors import CORSMiddleware
9import shutil
10
11class DialectClassifier(nn.Module):
12    def __init__(self, num_classes=6):
13        super(DialectClassifier, self).__init__()
14        self.conv_blocks = nn.Sequential(
15            nn.Conv2d(1, 32, kernel_size=(3, 3), padding='same'),
16            nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(kernel_size=(2, 2)),
17            nn.Conv2d(32, 64, kernel_size=(3, 3), padding='same'),
18            nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(kernel_size=(2, 2)),
19            nn.Conv2d(64, 128, kernel_size=(3, 3), padding='same'),
20            nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(kernel_size=(1, 2)),
21        )
22        self.lstm = nn.LSTM(input_size=128 * 20, hidden_size=128, num_layers=2, batch_first=True, bidirectional=True)
23        self.dropout = nn.Dropout(0.5)
24        self.fc1 = nn.Linear(128 * 2, 64)
25        self.fc_out = nn.Linear(64, num_classes)
26
27    def forward(self, x, input_lengths):
28        x = self.conv_blocks(x)
29        new_lengths = input_lengths // 8
30        new_lengths = torch.clamp(new_lengths, min=1)
31        max_len = x.size(3)
32        new_lengths = torch.clamp(new_lengths, max=max_len)
33        x = x.transpose(1, 3).contiguous()
34        B, T, H, C = x.shape
35        x = x.view(B, T, H * C)
36        x_packed = nn.utils.rnn.pack_padded_sequence(x, new_lengths.cpu(), batch_first=True, enforce_sorted=False)
37        x_packed, _ = self.lstm(x_packed)
38        x, _ = nn.utils.rnn.pad_packed_sequence(x_packed, batch_first=True)
39        x = torch.mean(x, dim=1)
40        x = self.dropout(x)
41        x = torch.relu(self.fc1(x))
42        return self.fc_out(x)
43
44
45app = FastAPI()
46
47
48app.add_middleware(
49    CORSMiddleware,
50    allow_origins=["*"],
51    allow_methods=["*"],
52    allow_headers=["*"],
53)
54
55device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
56model = DialectClassifier(num_classes=6).to(device)
57
58
59try:
60    model.load_state_dict(torch.load("best_dialect_classifier.pth", map_location=device))
61    model.eval()
62    print("Model loaded successfully!")
63except Exception as e:
64    print(f"Error loading model: {e}")
65
66# Mapping nhãn
67IDX_TO_CLS = { 
68    0: 'Giọng Bắc (North)', 
69    1: 'Bắc Trung Bộ (North Central Coast)',
70    2: 'Nam Trung Bộ (South Central Coast)', 
71    3: 'Tây Nguyên (Central Highland)',
72    4: 'Đông Nam Bộ (South East)',
73    5: 'Tây Nam Bộ (South West)',
74}
75
76# Transform
77valid_audio_transforms = nn.Sequential(
78    torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_mels=80),
79    torchaudio.transforms.AmplitudeToDB()
80)
81
82
83@app.post("/predict")
84async def predict_audio(file: UploadFile = File(...)):
85
86    temp_filename = f"temp_{file.filename}"
87    with open(temp_filename, "wb") as buffer:
88        shutil.copyfileobj(file.file, buffer)
89        
90    try:
91        audio, sr = librosa.load(temp_filename, sr=16000, mono=True)
92        audio_tensor = torch.from_numpy(audio).float()
93        
94        # Chuyển đổi audio thành spectrogram
95        spec = valid_audio_transforms(audio_tensor).squeeze(0).transpose(0, 1)
96        
97        input_length = torch.tensor([spec.shape[0]]).to(device)
98        
99        spec = spec.unsqueeze(0)
100        spec = spec.unsqueeze(1).transpose(2, 3)
101        spec = spec.to(device)
102        
103        # Dự đoán
104        with torch.no_grad():
105            output = model(spec, input_length)
106            probs = torch.softmax(output, dim=1)
107            score, predicted_idx = torch.max(probs, 1)
108            
109        result = {
110            "prediction": IDX_TO_CLS[predicted_idx.item()],
111            "confidence": float(score.item()),
112            "probabilities": {IDX_TO_CLS[i]: float(probs[0][i]) for i in range(6)}
113        }
114        
115        return result
116
117    except Exception as e:
118        return {"error": str(e)}
119    finally:
120        # Xóa file tạm
121        if os.path.exists(temp_filename):
122            os.remove(temp_filename)
123
124