Update app.py

app.py

@@ -6,43 +6,66 @@ from transformers import pipeline
 # Create a Flask app
 app = Flask(__name__)
 
-# Initialize the audio model at the start of the API
+# Initialize models at the start of the API
 audio_model = None
 
 def download_models():
     global audio_model
-    print("Downloading audio model...")
+    print("Downloading models...")
     # Download and load the audio model
     audio_model = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")
-    print("Audio model downloaded and ready to use.")
+    print("Model downloaded and ready to use.")
 
-# Download the model when the server starts
+# Download model when the server starts
 download_models()
 
 @app.route('/detect', methods=['POST'])
 def detect_deepfake():
+    data_type = request.form.get('type')  # "audio"
+    audio_file = request.files.get('audio_file')
     folder_path = request.form.get('folder_path')
 
-    if not folder_path or not os.path.isdir(folder_path):
-        return jsonify({"error": "Invalid folder path"}), 400
+    # If a single audio file is provided
+    if audio_file:
+        try:
+            # Save the uploaded file temporarily
+            file_path = os.path.join("/tmp", audio_file.filename)
+            audio_file.save(file_path)
 
-    results = {}
-    try:
-        # Process audio files only
-        for file_name in os.listdir(folder_path):
-            if file_name.endswith('.wav') or file_name.endswith('.mp3'):
-                file_path = os.path.join(folder_path, file_name)
-                result = audio_model(file_path)
-                results[file_name] = {item['label']: item['score'] for item in result}
+            # Perform detection
+            result = audio_model(file_path)
+            result_dict = {item['label']: item['score'] for item in result}
 
-        # Save results to a file
-        with open('detection_results.json', 'w') as f:
-            f.write(json.dumps(results))
+            # Remove the temporary file
+            os.remove(file_path)
 
-        return jsonify({"message": "Detection completed", "results": results}), 200
+            return jsonify({"message": "Detection completed", "results": result_dict}), 200
 
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        except Exception as e:
+            return jsonify({"error": str(e)}), 500
+
+    # If a folder path is provided
+    elif folder_path and os.path.isdir(folder_path):
+        results = {}
+        try:
+            for file_name in os.listdir(folder_path):
+                if file_name.endswith('.wav') or file_name.endswith('.mp3'):
+                    file_path = os.path.join(folder_path, file_name)
+                    result = audio_model(file_path)
+                    results[file_name] = {item['label']: item['score'] for item in result}
+
+            # Save results to a file
+            with open('detection_results.json', 'w') as f:
+                f.write(json.dumps(results))
+
+            return jsonify({"message": "Detection completed", "results": results}), 200
+
+        except Exception as e:
+            return jsonify({"error": str(e)}), 500
+
+    # Invalid request if neither audio file nor valid folder path is provided
+    else:
+        return jsonify({"error": "Invalid input. Provide an audio file or a valid folder path."}), 400
 
 if __name__ == '__main__':
     # Run the Flask app