from flask import Flask, render_template, request
import tensorflow as tf
import numpy as np
from keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from io import BytesIO
import base64

app = Flask(__name__)

# Load the model
incept_model = load_model('best_model_2.h5')
IMAGE_SHAPE = (224, 224)
classes = ['benign', 'malignant', 'normal']

# Function to prepare the image
def prepare_image(file):
    img = load_img(BytesIO(file.read()), target_size=IMAGE_SHAPE)
    img_array = img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    return tf.keras.applications.efficientnet.preprocess_input(img_array)

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

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return redirect(request.url)
    file = request.files['file']
    if file.filename == '':
        return redirect(request.url)

    # Prepare the image for prediction
    img = prepare_image(file)
    res = incept_model.predict(img)
    pred = classes[np.argmax(res)]

    # Encode image to display in the result page
    file.seek(0)  # Reset file pointer to the beginning
    img_bytes = file.read()
    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
    img_data = f"data:image/jpeg;base64,{img_base64}"

    return render_template('result.html', prediction=pred, image_data=img_data)

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