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



app = Flask(__name__)
CORS(app)

#dataset clasess
class_name = {
    0: 'Blister Blight',
    1: 'Brown Blight',
    2: 'Gray Blight',
    3: 'Healthy',
    4: 'White Spot'
}

#load saved model
model_dir = './CNN_TEA_MODEL.h5'
model = load_model(model_dir)


    
        

@app.route("/predict", methods=["POST"])
def predictTest():
    
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if file:
        # Convert the file storage to PIL Image and ensure it's in RGB
        img = Image.open(io.BytesIO(file.read())).convert('RGB')  # Added .convert('RGB')
        img = img.resize((256, 256))
        img_array = np.array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = img_array / 255.0  # Normalize
        
        predictions = model.predict(img_array)
        #get the class with the highest probability
        predicted_class = np.argmax(predictions, axis=1)
        predicted_class_name = class_name[predicted_class[0]]
        
        result = {"class": predicted_class_name}
        print("Prediction: ", result)
        print(predictions)
        
        return jsonify({'prediction': result})






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