Update app.py

app.py

@@ -1,51 +1,55 @@
-import librosa
-import numpy as np
-import tensorflow as tf
-
-def load_audio(file_path):
-    audio, sr = librosa.load(file_path, sr=16000)
-    return audio, sr
-
-def extract_features(audio, sr):
-    mfccs = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=13)
-    return mfccs
-
-# Load pre-trained model (example: DeepSpeech)
-model = tf.keras.models.load_model('deepspeech_model.h5')
-
-def speech_to_text(audio):
-    features = extract_features(audio, 16000)
-    features = np.expand_dims(features, axis=0)  # Model expects batch dimension
-    text = model.predict(features)
-    return text
-
-from transformers import MarianMTModel, MarianTokenizer
-
-# Load pre-trained translation model
-model_name = 'Helsinki-NLP/opus-mt-en-hi'
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-translation_model = MarianMTModel.from_pretrained(model_name)
-
-def translate_text(text, tokenizer, model):
-    inputs = tokenizer(text, return_tensors="pt", padding=True)
-    translated = model.generate(**inputs)
-    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
-    return translated_text
-
-import datetime
-
-def process_audio_for_translation(audio_file_path):
-    current_time = datetime.datetime.now().time()
-    if current_time < datetime.time(18, 0):
-        return "Translation service is available only after 6 PM IST"
-
-    audio, sr = load_audio(audio_file_path)
-    english_text = speech_to_text(audio)
-    hindi_text = translate_text(english_text, tokenizer, translation_model)
-    return hindi_text
-
-# Example usage
-audio_file_path = 'path_to_audio_file.wav'
-translated_text = process_audio_for_translation(audio_file_path)
-print(translated_text)
+from flask import Flask, request, jsonify
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer, pipeline
+from pydub import AudioSegment
+import torch
+import torchaudio
+from datetime import datetime, time
+import pytz
+
+app = Flask(__name__)
+
+# Load speech recognition model and tokenizer
+tokenizer = Wav2Vec2Tokenizer.from_pretrained("facebook/wav2vec2-base-960h")
+model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+
+# Load translation pipeline
+translation_pipeline = pipeline("translation_en_to_hi")
+
+# Function to preprocess audio
+def preprocess_audio(audio_file):
+    audio = AudioSegment.from_file(audio_file)
+    audio = audio.set_frame_rate(16000)
+    audio.export("processed.wav", format="wav")
+    waveform, sample_rate = torchaudio.load("processed.wav")
+    return waveform
+
+# Function to check if the current time is after 6 PM IST
+def is_after_6pm_ist():
+    ist = pytz.timezone('Asia/Kolkata')
+    current_time = datetime.now(ist).time()
+    return current_time >= time(18, 0)
+
+@app.route('/translate', methods=['POST'])
+def translate():
+    if not is_after_6pm_ist():
+        return jsonify({"error": "Service is available only after 6 PM IST"}), 403
+
+    if 'audio' not in request.files:
+        return jsonify({"error": "No audio file provided"}), 400
+
+    audio_file = request.files['audio']
+    waveform = preprocess_audio(audio_file)
+    
+    input_values = tokenizer(waveform.squeeze().numpy(), return_tensors="pt").input_values
+    logits = model(input_values).logits
+    predicted_ids = torch.argmax(logits, dim=-1)
+    transcription = tokenizer.batch_decode(predicted_ids)[0]
+
+    translation = translation_pipeline(transcription)
+    translated_text = translation[0]['translation_text']
+
+    return jsonify({"transcription": transcription, "translation": translated_text})
+
+if __name__ == '__main__':
+    app.run(debug=True, host='0.0.0.0', port=8080)