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

# Initialize Flask application
app = Flask(__name__)

# Load your trained model
model = load_model('./butterfly_classifier.h5')

# Define a function to preprocess an image
def preprocess_image(image_path):
    img = image.load_img(image_path, target_size=(224, 224))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Define a route to predict butterfly species
@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'})
    
    file = request.files['file']
    
    # Save the uploaded file to a temporary location
    file_path = 'temp.jpg'
    file.save(file_path)
    
    # Preprocess the image
    processed_image = preprocess_image(file_path)
    
    # Make prediction
    prediction = model.predict(processed_image)
    
    # Get predicted label (assuming your classes are encoded as integers)
    predicted_class = np.argmax(prediction, axis=1)
    
    # Clean up: remove the temporary file
    os.remove(file_path)
    
    # Return the result as JSON
    return jsonify({'predicted_class': predicted_class.item()})

# Define a welcome route
@app.route('/')
def welcome():
    return 'Welcome to Butterfly Classification API'

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