from flask import Flask, request, jsonify
from fastai.vision.all import *

app = Flask(__name__)
learn = load_learner("export.pkl")
labels = learn.dls.vocab

ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "gif"}
def allowed_file(filename):
  return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route("/api/predict", methods=["POST"])
def predict():
  if "file" not in request.files:
    return jsonify({"message": "No file part in the request"}), 400
  file = request.files["file"]

  if file.filename == "":
    return jsonify({"message": "No selected file"}), 400

  if file and allowed_file(file.filename):
    img = PILImage.create(file)
    pred, pred_idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}
  else:
    return jsonify({"message": "File type not allowed"}), 400

if __name__ == "__main__":
  app.run(debug=True)