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

app = Flask(__name__)

# Load your model
model = load_model('./best_model.h5', compile=False)

# Itsekiri digit labels
itsekiri_labels = [
    'Méene', 'Méji', 'Métà', 'Mérin', 'Márùn',
    'Méfa', 'Méje', 'Méjọ', 'Méṣan', 'Méwà'
]

IMG_SIZE = (228, 228)

def preprocess_image(base64_str):
    header, encoded = base64_str.split(',', 1)
    img_bytes = base64.b64decode(encoded)
    img = Image.open(BytesIO(img_bytes)).convert('RGB')
    img = img.resize(IMG_SIZE)
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array.astype('float32') / 255.0
    return img_array

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json(force=True)
    img_data = data['image']
    img_array = preprocess_image(img_data)
    pred_probs = model.predict(img_array)
    pred_index = np.argmax(pred_probs)
    pred_label = itsekiri_labels[pred_index]
    confidence = float(pred_probs[0][pred_index] * 100)
    return jsonify({
        'prediction': pred_label,
        'confidence': confidence
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)