Update api.py

api.py

@@ -1,11 +1,11 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, File, UploadFile
 from huggingface_hub import from_pretrained_keras
 import numpy as np
 from PIL import Image
 import io
 import tensorflow as tf
 
-app = Flask(__name__)
+app = FastAPI()
 
 # Load the model
 # model = from_pretrained_keras("MissingBreath/recycle-garbage-model")
@@ -14,21 +14,21 @@ model = tf.keras.models.load_model('_9217')
 # Class labels
 # class_labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
-@app.route('/classify', methods=['POST'])
-def classify():
-    file = request.files['image']
-    if file:
-        img = Image.open(io.BytesIO(file.read()))
+@app.post("/classify")
+async def classify(image: UploadFile = File(...)):
+    if image is not None:
+        img = Image.open(io.BytesIO(await image.read()))
         img = img.resize((128, 128))
         img_array = np.array(img) / 255.0
         img_array = np.expand_dims(img_array, axis=0)
         predictions = model.predict(img_array)
         predicted_class_idx = np.argmax(predictions)
         # predicted_class = class_labels[predicted_class_idx]
-        # return jsonify({'prediction': predicted_class})
-        return jsonify({'prediction': predicted_class_idx})
+        # return {"prediction": predicted_class}
+        return {"prediction": predicted_class_idx}
     else:
-        return jsonify({'error': 'No image provided'}), 400
+        return {"error": "No image provided"}
 
-if __name__ == '__main__':
-    app.run(debug=False)
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run(app, host="0.0.0.0", port=8000)