import numpy as np
import tensorflow as tf
from flask import Flask, request, jsonify, make_response
from PIL import Image
import io

app = Flask(__name__)

classes = ['Apple___Apple_scab',
 'Apple___Black_rot',
 'Apple___Cedar_apple_rust',
 'Apple___healthy',
 'Blueberry___healthy',
 'Cherry_(including_sour)___Powdery_mildew',
 'Cherry_(including_sour)___healthy',
 'Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot',
 'Corn_(maize)___Common_rust_',
 'Corn_(maize)___Northern_Leaf_Blight',
 'Corn_(maize)___healthy',
 'Grape___Black_rot',
 'Grape___Esca_(Black_Measles)',
 'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)',
 'Grape___healthy',
 'Orange___Haunglongbing_(Citrus_greening)',
 'Peach___Bacterial_spot',
 'Peach___healthy',
 'Pepper,_bell___Bacterial_spot',
 'Pepper,_bell___healthy',
 'Potato___Early_blight',
 'Potato___Late_blight',
 'Potato___healthy',
 'Raspberry___healthy',
 'Soybean___healthy',
 'Squash___Powdery_mildew',
 'Strawberry___Leaf_scorch',
 'Strawberry___healthy',
 'Tomato___Bacterial_spot',
 'Tomato___Early_blight',
 'Tomato___Late_blight',
 'Tomato___Leaf_Mold',
 'Tomato___Septoria_leaf_spot',
 'Tomato___Spider_mites Two-spotted_spider_mite',
 'Tomato___Target_Spot',
 'Tomato___Tomato_Yellow_Leaf_Curl_Virus',
 'Tomato___Tomato_mosaic_virus',
 'Tomato___healthy']

from tensorflow.keras.models import load_model


# Load the model without loading the weights
model = load_model('model.h5', compile=False)
for layer in model.layers:
    if isinstance(layer, tf.keras.layers.BatchNormalization):
        # Modify the configuration of BatchNormalization layer
        config = layer.get_config()
        config['axis'] = 3  # Ensure axis is set correctly as an integer
        # Update the layer with the modified configuration
        new_layer = tf.keras.layers.BatchNormalization.from_config(config)
        # Replace the old layer with the updated one
        model.layers[model.layers.index(layer)] = new_layer
model.save('modified_model.h5')



# custom_objects = {'BatchNormalization': tf.keras.layers.BatchNormalization}
loaded_model = load_model('modified_model.h5')

# Load the h5 model
# model = tf.keras.models.load_model("model.h5")

# Preprocess the image
def preprocess_image(image, target_size):
    image = image.convert('RGB')
    image = image.resize(target_size)
    image = np.array(image, dtype=np.float32)
    image = np.expand_dims(image, axis=0)
    image = image / 255.0  # Normalize if required
    return image

@app.after_request
def after_request(response):
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
    response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
    return response

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400
    
    file = request.files['file']
    
    if file.filename == '':
        return jsonify({"error": "No file selected for uploading"}), 400
    
    if file:
        # Read the image
        image = Image.open(io.BytesIO(file.read()))

        # Preprocess the image
        target_size = (224, 224)
        image = preprocess_image(image, target_size)

        # Make prediction
        predictions = loaded_model.predict(image)
        predicted_class = np.argmax(predictions, axis=-1)[0]
        confidence = np.max(predictions, axis=-1)[0]

        response = jsonify({
            "predicted_class": classes[int(predicted_class)],
            "confidence": float(confidence)
        })
        response.headers.add('Access-Control-Allow-Origin', '*')
        return response

    return jsonify({"error": "An error occurred during prediction"}), 500

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