Update app.py

app.py

@@ -1,15 +1,16 @@
 from flask import Flask, jsonify, request
 from flask_cors import CORS
-import numpy as np
-import matplotlib.pyplot as plt
-import base64
-import io
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.models as models
-from PIL import Image
 from torchvision import transforms
+import torchaudio
+import numpy as np
+import matplotlib.pyplot as plt
+import base64
+import io
+from PIL import Image
 from ultralytics import YOLO
 
 
@@ -21,6 +22,9 @@ CORS(app)
 idx_to_class_resnet50 = {0 : "Genuine" , 1:'Printed Paper' , 2 : 'Replayed'}
 idx_to_class_yolo9 = idx_to_class_yolo9 = {0: 'Genuine', 1: 'Printed Paper', 2: 'Replayed', 3: 'Paper Mask'}
 idx_to_class_resnet50_celeba = {0 : "Genuine" , 1:'Printed Paper' , 2 : 'Paper Cut',3:'Replayed',4:'3D Mask'}
+binary_labels = ['real','spoof']
+
+
 transform_data_resnet50=transforms.Compose([
     transforms.Resize(size=(224,224)),
     transforms.ToTensor()
@@ -31,6 +35,28 @@ transform_data_resnet50_celeba=transforms.Compose([
     transforms.Resize((224,224), antialias=True)
 ])
 
+def process_audio(encoded_audio):
+    decoded_audio = base64.b64decode(encoded_audio)
+    audio_bytes = io.BytesIO(decoded_audio)
+    waveform, sample_rate = torchaudio.load(audio_bytes)
+    
+    if waveform.size(0) > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)  # Convert to mono by averaging channels
+    
+    mel_spectrogram = torchaudio.transforms.MelSpectrogram(n_mels=80)(waveform).squeeze(0)
+    num_frames = mel_spectrogram.size(1)
+    target_length = 400
+    
+    if num_frames < target_length:
+        padding = target_length - num_frames
+        mel_spectrogram = torch.cat([mel_spectrogram, torch.zeros(mel_spectrogram.size(0), padding)], dim=1)
+    else:
+        mel_spectrogram = mel_spectrogram[:, :target_length]
+    
+    mel_spectrogram = mel_spectrogram.transpose(0, 1)
+    length = torch.tensor([mel_spectrogram.size(0)])
+    return mel_spectrogram.unsqueeze(0) ,length
+
 model_resnet50 = models.resnet50(weights=False)
 num_classes = 3 
 model_resnet50.fc = nn.Linear(model_resnet50.fc.in_features, num_classes)
@@ -44,12 +70,55 @@ model_resnet50_celeba.load_state_dict(torch.load('resnet50_model_weights_celeba.
 model_resnet50_celeba.eval()
 
 model_yolo9 = YOLO('yolo9_best.pt')
+
+
+
+class ConformerClassifier(torch.nn.Module):
+    def __init__(self, input_dim, num_classes, num_heads, ffn_dim, num_layers, depthwise_conv_kernel_size,dropout=0.0,use_group_norm=False,convolution_first=False):
+        super(ConformerClassifier, self).__init__()
+        self.conformer = torchaudio.models.Conformer(
+            input_dim=input_dim,
+            num_heads=num_heads,
+            ffn_dim=ffn_dim,
+            num_layers=num_layers,
+            depthwise_conv_kernel_size=depthwise_conv_kernel_size,
+            dropout=dropout,
+            use_group_norm=use_group_norm,
+            convolution_first=convolution_first
+        )
+        self.fc = torch.nn.Linear(input_dim, num_classes)
+    
+    def forward(self, x, lengths):
+        x,length = self.conformer(x, lengths)
+        x = x.mean(dim=1)
+        x = self.fc(x)
+        return x
+    
+voice_binary_model = ConformerClassifier(
+    input_dim=80,
+    num_classes=2,  
+    num_heads=4,
+    ffn_dim=128,
+    num_layers=4,
+    depthwise_conv_kernel_size=7,
+    dropout=0.3,
+    use_group_norm=False,
+    convolution_first=True   
+)
+voice_binary_model.load_state_dict(torch.load('binary_voice_model.pth',map_location='cpu'))
+voice_binary_model.eval()
+
 print('Models Loaded Successfully')
 
 
 
 
-@app.route('/', methods=['GET'])
+@app.route('/')
+def home():
+    return "Welcome to the Flask API!"
+
+
+@app.route('/api/data', methods=['GET'])
 def get_data():
     img = plt.imread('test1.jpeg')
     img_arr = np.array(img)
@@ -65,11 +134,6 @@ def get_data():
     }
     return jsonify(data)
 
-
-@app.route('/test')
-def home():
-    return "Welcome to the Flask API!"
-
 @app.route('/', methods=['POST'])
 def post_data():
     try:
@@ -146,3 +210,20 @@ def post_test_data():
     return jsonify(response), 201
 
 
+@app.route('/api/voice', methods=['POST'])
+def post_api_voice():
+    data = request.json
+    encoded_audio = data['base64']
+    mel_spectrogram, length = process_audio(encoded_audio)
+    with torch.no_grad():
+        output = voice_binary_model(mel_spectrogram,length)
+        prob  = torch.nn.functional.softmax(output[0], dim=0)
+        pred = torch.argmax(prob).item()
+        category = binary_labels[pred]
+        probs_dict = {binary_labels[i]:prob[i] for i in range(len(binary_labels))}
+    response = {
+        'message': 'Data received!',
+        'class' : category,
+        'probs' : probs_dict
+    }
+    return jsonify(response), 201