Update main.py

main.py

@@ -5,13 +5,13 @@ from flask import Flask, request, jsonify
 import torch
 import torch.nn.functional as F
 from facenet_pytorch import MTCNN, InceptionResnetV1
-import os
 import numpy as np
 from PIL import Image
 import cv2
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from pytorch_grad_cam.utils.image import show_cam_on_image
+from io import BytesIO
 
 app = Flask(__name__)
 
@@ -39,23 +39,11 @@ model.load_state_dict(checkpoint['model_state_dict'])
 model.to(DEVICE)
 model.eval()
 
-UPLOAD_FOLDER = os.getcwd()  # Save the uploaded image to the current working directory
-
-def predict(input_image):
+def predict(image_bytes):
     """Predict the label of the input_image"""
     true_label = request.form['true_label']
-    if 'file' not in request.files:
-        return jsonify({"error": "No file part"}), 400
-    file = request.files['file']
-    if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
-    
-    # Save the uploaded image to the uploads folder
-    input_image_path = os.path.join(UPLOAD_FOLDER, file.filename)
-    file.save(input_image_path)
-    
-    input_image = Image.open(input_image_path)
-    face = mtcnn(input_image)
+    image = Image.open(BytesIO(image_bytes))
+    face = mtcnn(image)
     if face is None:
         return jsonify({"error": "No face detected"}), 400
 
@@ -99,7 +87,12 @@ def index():
 
 @app.route("/predict", methods=['POST'])
 def predict_endpoint():
-    return predict(request)
+    if 'file' not in request.files:
+        return jsonify({"error": "No file part"}), 400
+    file = request.files['file']
+    if file.filename == '':
+        return jsonify({"error": "No selected file"}), 400
+    return predict(file.read())
 
 if __name__ == "__main__":
     app.run(debug=True, port=7860)  # Change the port to 7860