from flask import Flask, request, jsonify
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing import image

app = Flask(__name__)
model = tf.keras.models.load_model("MobileNet_Fire.h5")

class_labels = {0: "Fake", 1: "Low", 2: "Medium", 3: "High"}  # Update as per your training

@app.route("/predict", methods=["POST"])
def predict():
    file = request.files["file"]
    img = image.load_img(file, target_size=(128, 128))
    img_array = image.img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    predictions = model.predict(img_array)
    predicted_class = class_labels[np.argmax(predictions)]
    confidence = float(np.max(predictions))

    return jsonify({"prediction": predicted_class, "confidence": confidence})

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