API / app.py
ParthBarot's picture
Update app.py
7ec95bc verified
raw
history blame
2.03 kB
import os
import json
from flask import Flask, jsonify, request
from transformers import pipeline
from pydub import AudioSegment
# Create a Flask app
app = Flask(__name__)
# Initialize models at the start of the API
audio_model = None
def download_models():
global audio_model
print("Downloading models...")
# Download and load the audio model with padding enabled
audio_model = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2", padding=True)
print("Model downloaded and ready to use.")
# Download model when the server starts
download_models()
def convert_audio_to_wav(input_path, output_path):
# Convert any audio format to WAV using pydub
audio = AudioSegment.from_file(input_path)
audio.export(output_path, format="wav")
@app.route('/detect', methods=['POST'])
def detect_deepfake():
# Expect an audio file in the request
audio_file = request.files.get('audio_file')
# If a single audio file is provided
if audio_file:
try:
# Save the uploaded file temporarily
input_path = os.path.join("/tmp", audio_file.filename)
audio_file.save(input_path)
# Convert the file to WAV format
output_path = os.path.splitext(input_path)[0] + '.wav'
convert_audio_to_wav(input_path, output_path)
# Perform detection
result = audio_model(output_path)
result_dict = {item['label']: item['score'] for item in result}
# Remove the temporary files
os.remove(input_path)
os.remove(output_path)
return jsonify({"message": "Detection completed", "results": result_dict}), 200
except Exception as e:
return jsonify({"error": str(e)}), 500
# Invalid request if no audio file is provided
else:
return jsonify({"error": "Invalid input. Please provide an audio file."}), 400
if __name__ == '__main__':
# Run the Flask app
app.run(host='0.0.0.0', port=7860)