import argparse

from PIL import Image
from flask import Flask, request, jsonify
from pathlib import Path

app = Flask(__name__)


@app.before_first_request
def init():
    from fb_detr.predictor import Predictor

    global PRED

    PRED = Predictor(args.resume)


@app.route("/", methods=["GET", "POST"])
def predict_request():
    if request.method == "POST":
        image_folder_path = request.form.get("image_folder_path")
        images = list(map(Image.open, Path(image_folder_path).glob("*.png")))
        results = PRED.predict(images, format_output=True)
        for result in results:
            return jsonify(result)
    if request.method == "GET":
        return "Not implemented"


parser = argparse.ArgumentParser()
parser.add_argument("--resume", required=True)
args = parser.parse_args()

app.run()