is_it_a_panda / server.py
timothyoei's picture
Added support for local flask server
11e2e05
raw
history blame contribute delete
911 Bytes
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)