kingav/bts

1from flask import Flask, request, render_template, redirect, url_for
2import numpy as np
3import os
4from tensorflow.keras.models import load_model
5from werkzeug.utils import secure_filename
6import cv2
7import matplotlib.pyplot as plt
8from io import BytesIO
9import base64
10import math
11
12app = Flask(__name__)
13app.config['UPLOAD_FOLDER'] = 'uploads/'
14
15# Load the trained model
16model = load_model('model.h5')
17
18# Ensure the upload folder exists
19os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
20
21def preprocess_image(image_path):
22    image = np.load(image_path)
23    if image.ndim == 3:
24        image = np.expand_dims(image, axis=-1)
25    target_shape = (64, 64, 64, 1)
26    image = np.pad(image, ((0, max(0, target_shape[0] - image.shape[0])),
27                           (0, max(0, target_shape[1] - image.shape[1])),
28                           (0, max(0, target_shape[2] - image.shape[2])),
29                           (0, 0)), mode='constant')
30    image = image[:target_shape[0], :target_shape[1], :target_shape[2], :]
31    return np.expand_dims(image, axis=0)
32
33@app.route('/', methods=['GET', 'POST'])
34def upload_file():
35    if request.method == 'POST':
36        if 'file' not in request.files:
37            return redirect(request.url)
38        file = request.files['file']
39        if file.filename == '':
40            return redirect(request.url)
41        if file:
42            filename = secure_filename(file.filename)
43            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
44            file.save(file_path)
45            return redirect(url_for('predict', filename=filename))
46    return render_template('upload.html')
47
48@app.route('/predict/<filename>')
49def predict(filename):
50    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
51    image = preprocess_image(file_path)
52    prediction = model.predict(image)
53    
54    # Flatten the prediction
55    # flattened_prediction = flatten_nested_list(prediction)
56    prediction_array = np.array(prediction)
57    
58    # Reshape the prediction to a 2D image
59    side_length = int(math.sqrt(prediction_array.size))
60    binary_prediction = (prediction_array > 0.5).astype(np.uint8).reshape(side_length, -1)
61    
62    print("---------- predicted ----------")
63    print("Prediction shape:", binary_prediction.shape)
64
65    # Save the prediction mask
66    mask_path = os.path.join(app.config['UPLOAD_FOLDER'], 'mask_' + filename)
67    np.save(mask_path, binary_prediction)
68
69    # Create a plot of the prediction
70    plt.figure(figsize=(10, 10))
71    plt.imshow(binary_prediction, cmap='gray')
72    plt.axis('off')
73    plt.title('Prediction')
74    
75    # Save the plot to a BytesIO object
76    img_buffer = BytesIO()
77    plt.savefig(img_buffer, format='png', bbox_inches='tight')
78    img_buffer.seek(0)
79    img_str = base64.b64encode(img_buffer.getvalue()).decode()
80    plt.close()
81
82    return render_template('result.html', original_image=filename, mask_image='mask_' + filename, prediction_image=img_str)
83
84def flatten_nested_list(nested_list):
85    flattened = []
86    for item in nested_list:
87        if isinstance(item, (list, np.ndarray)):
88            flattened.extend(flatten_nested_list(item))
89        else:
90            flattened.append(item)
91    return flattened
92
93if __name__ == '__main__':
94    app.run(debug=True)