from flask import Flask, request, jsonify
import tensorflow as tf
import numpy as np
import cv2
import base64

app = Flask(__name__)

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

# Function to decode base64 image
def decode_image(image_data):
    image_bytes = base64.b64decode(image_data)
    image_np = np.frombuffer(image_bytes, dtype=np.uint8)
    image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
    image = cv2.resize(image, (224, 224))  # Adjust based on your model
    image = image / 255.0  # Normalize if needed
    return image.reshape(1, 224, 224, 3)

# API endpoint for prediction
@app.route('/predict', methods=['POST'])
def predict():
    try:
        data = request.json['image']
        image = decode_image(data)
        prediction = model.predict(image).tolist()
        return jsonify({'prediction': prediction})
    except Exception as e:
        return jsonify({'error': str(e)})

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