from flask import Flask, request, jsonify
from tensorflow.keras.preprocessing import image
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image
from io import BytesIO

app = Flask(__name__)

# Load the trained models
# model_female_leg = load_model('trained_model_female_leg.h5')
# model_female_arm = load_model('trained_model_female_arm.h5')
# model_male_arm = load_model('trained_model_male_arm.h5')
#model_male_leg = load_model('YourModelName.h5')
#model_male_arm = load_model('YourModelName.h5')
#model_female_arm = load_model('YourModelName.h5')
#model_female_leg = load_model('model.h5')
#model_male_leg = load_model('trained_model_male_leg.h5')


model_male_leg = load_model('model.h5')
model_male_arm = load_model('model.h5')
model_female_leg = load_model('YourModelName.h5')
model_female_arm = load_model('model.h5')

# Define class labels for each model
class_label_male_leg = ['High', 'Moderate', 'Low']
class_label_male_arm = ['High', 'Moderate', 'Low']
class_labels_female_leg = ['High', 'Moderate', 'Low']
class_labels_female_arm = ['High', 'Moderate', 'Low']

# Define route for model prediction for model 1
@app.route('/predict_model_male_leg', methods=['POST'])
def predict_model1():
    return predict(request.files['file'], model_male_leg, class_label_male_leg)

# Define route for model prediction for model 2
@app.route('/predict_model_male_arm', methods=['POST'])
def predict_model2():
    return predict(request.files['file'], model_male_arm, class_label_male_arm)

# Define route for model prediction for model 3
@app.route('/predict_model_female_leg', methods=['POST'])
def predict_model3():
    return predict(request.files['file'], model_female_leg, class_labels_female_leg)

# Define route for model prediction for model 4
@app.route('/predict_model_female_arm', methods=['POST'])
def predict_model4():
    return predict(request.files['file'], model_female_arm, class_labels_female_arm)

# Define route for ping
@app.route('/', methods=['GET'])
def ping():
    return jsonify({'PING': 'Success!'})

def predict(file, model, class_labels):
    # Check if file is provided
    if not file:
        return jsonify({'error': 'No file provided'})

    # Load and preprocess the image
    img = Image.open(BytesIO(file.read()))  # Convert FileStorage to io.BytesIO
    img = img.resize((150, 150))  # Resize image to match model's input shape
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array /= 255.0

    # Make prediction
    prediction = model.predict(img_array)

    # Interpret the result
    predicted_class = np.argmax(prediction)
    predicted_label = class_labels[predicted_class]

    return jsonify( predicted_label)

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