app.py

from flask import Flask, render_template, request, jsonify, send_from_directory
import os
from werkzeug.utils import secure_filename
import numpy as np
from tensorflow.keras.models import load_model
import datetime
from pytz import timezone

app = Flask(__name__)

# Define the folder for predictions
PREDICTIONS_FOLDER = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'predictions')
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'inputs')
app.config['PREDICTIONS_FOLDER'] = PREDICTIONS_FOLDER
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Ensure the predictions folder exists
os.makedirs(PREDICTIONS_FOLDER, exist_ok=True)
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

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

# Set Current Date and Time for naming the files
now = datetime.datetime.now(timezone('Europe/Berlin'))
current_time_label = now.strftime("%H%M")
current_date_label = now.strftime("%Y%m%d_")
TIMESTAMP = current_date_label+current_time_label

# Flag to indicate whether predictions are available
predictions_available = False

@app.route('/')
def index():
    global predictions_available
    return render_template('index.html', predictions_available=predictions_available)

@app.route('/predict', methods=['POST'])
def predict():

    global predictions_available
    
    from matplotlib.pyplot import imsave
    
    # Get input data from the request
    inputfile = request.files['input_file']
    input_data = np.load(inputfile.stream)
    randomint = np.random.randint(0, input_data.shape[0])
    input_filename = secure_filename(inputfile.filename)
    np.save(os.path.join(UPLOAD_FOLDER, f"{os.path.splitext(input_filename)[0]}" + "_" + TIMESTAMP + ".npy"), input_data)
    input_figname = f"{os.path.splitext(input_filename)[0]}" + "_" + TIMESTAMP + ".png"
    imsave(os.path.join(UPLOAD_FOLDER, input_figname), input_data[randomint, :, :, 0], cmap='gray')

    # Perform prediction
    TotalData_mean = -595.6387043163428
    TotalData_std = 627.3217112391814
    input_data = (input_data - TotalData_mean) / TotalData_std
    predictions = model.predict(input_data)
    predictions = predictions*TotalData_std + TotalData_mean

    # Save predictions as an npy file
    output_filename = f"{os.path.splitext(input_filename)[0]}_pred" + "_" + TIMESTAMP + ".npy"
    output_figname = f"{os.path.splitext(input_filename)[0]}_pred" + "_" + TIMESTAMP + ".png"
    np.save(os.path.join(PREDICTIONS_FOLDER, output_filename), predictions)
    imsave(os.path.join(PREDICTIONS_FOLDER, output_figname), predictions[randomint, :, :, 0], cmap='gray')

    # Set the flag to indicate predictions are available
    predictions_available = True

    # Redirect to the home page after prediction
    return render_template('index.html', 
                           predictions_available=predictions_available, 
                           download_filename=output_filename, 
                           input_fig=input_figname,
                           output_fig=output_figname)

@app.route('/download/<filename>')
def download_file(filename):
    file_path = os.path.join(PREDICTIONS_FOLDER, filename)
    print(f"Trying to access file: {file_path}")
    return send_from_directory(PREDICTIONS_FOLDER, filename, as_attachment=True)

@app.route('/inputs/<filename>')
def input_image(filename):
    return send_from_directory(UPLOAD_FOLDER, filename, as_attachment=True)

@app.route('/predictions/<filename>')
def predictions_image(filename):
    return send_from_directory(PREDICTIONS_FOLDER, filename, as_attachment=True)


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