from flask import Flask, request, jsonify
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import  load_img
import numpy as np
import os

app = Flask(__name__)


@app.route('/')
def indexfn():
    return "<h1> Welcome </h1>"

@app.route('/predict', methods = ['POST'])
def predictfn():
    class_mask = ['Staircase not detected', 'Staircase Detected']
    file = request.files['file']
    image_path = "./images/" + file.filename
    file.save(image_path)

    try:
        img = load_img(image_path, target_size=(224,224))
        img = np.asarray(img)
        img = np.expand_dims(img, axis=0)
        saved_model = load_model('vgg16_1.h5')
        output = saved_model.predict(img)
        pred = class_mask[np.argmax(output)]
        os.remove(path=image_path)
        return jsonify({'type': pred})

    except BaseException as err:
        return jsonify({'error': 'error at prediction'})

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