kingav/bts
0
1import gradio as gr2import numpy as np3import os4from tensorflow.keras.models import load_model5import matplotlib.pyplot as plt6from io import BytesIO7import base648import shutil9 10# Load the trained model11model = load_model('model.h5')12 13# Configure upload and static folders14UPLOAD_FOLDER = 'uploads/'15STATIC_FOLDER = os.path.join('static', 'uploads')16 17# Ensure the upload and static folders exist18os.makedirs(UPLOAD_FOLDER, exist_ok=True)19os.makedirs(STATIC_FOLDER, exist_ok=True)20 21def combine_slices(image, method='grid'):22 """23 Combine slices from a 3D or 4D image into a single 2D image.24 """25 # Handle 4D arrays (e.g., multiple channels or batches)26 if image.ndim == 4:27 if image.shape[-1] == 1:28 image = image[..., 0] # Remove single channel dimension29 else:30 image = np.mean(image, axis=-1) # Average over channels/batch31 32 # Now, the image should be 3D (height, width, depth)33 if image.ndim != 3:34 raise ValueError(f"Expected a 3D array after processing, but got shape: {image.shape}")35 36 height, width, depth = image.shape37 38 if method == 'average':39 # Average across the depth axis40 return np.mean(image, axis=2)41 42 elif method == 'grid':43 # Concatenate slices into a grid44 grid_size = int(np.ceil(np.sqrt(depth)))45 46 # Create an empty array to hold the grid, with appropriate padding if necessary47 grid_image = np.zeros((grid_size * height, grid_size * width))48 49 # Fill the grid with slices50 for i in range(depth):51 row = i // grid_size52 col = i % grid_size53 grid_image[row*height:(row+1)*height, col*width:(col+1)*width] = image[:, :, i]54 55 return grid_image56 57 else:58 raise ValueError("Invalid method. Choose from 'average' or 'grid'.")59 60def preprocess_image(image_path):61 image = np.load(image_path)62 if image.ndim == 3:63 image = np.expand_dims(image, axis=-1)64 target_shape = (64, 64, 64, 1)65 image = np.pad(image, ((0, max(0, target_shape[0] - image.shape[0])),66 (0, max(0, target_shape[1] - image.shape[1])),67 (0, max(0, target_shape[2] - image.shape[2])),68 (0, 0)), mode='constant')69 image = image[:target_shape[0], :target_shape[1], :target_shape[2], :]70 return np.expand_dims(image, axis=0)71 72def copy_to_static(filename):73 """Copy uploaded file to static folder for display"""74 src = os.path.join(UPLOAD_FOLDER, filename)75 dst = os.path.join(STATIC_FOLDER, filename)76 shutil.copy(src, dst)77 78# def predict_and_visualize(file_obj):79# # Save the uploaded file to the upload folder80# filename = secure_filename(file_obj.name)81# file_path = os.path.join(UPLOAD_FOLDER, filename)82# file_obj.save(file_path)83# copy_to_static(filename) # Copy file to static folder84 85# # Preprocess the image86# image = preprocess_image(file_path)87# prediction = model.predict(image)88 89# # Handle the model output90# bridge_output, logits = prediction91 92# # Use the logits (second output) for visualization93# binary_prediction = (logits[0, ..., 0] > 0.5).astype(np.uint8)94 95# print("---------- predicted ----------")96# print("Prediction shape:", binary_prediction.shape)97 98# # Save the prediction mask99# mask_filename = 'mask_' + filename100# mask_path = os.path.join(UPLOAD_FOLDER, mask_filename)101# np.save(mask_path, binary_prediction)102# copy_to_static(mask_filename) # Copy mask to static folder103 104# # Visualize original image and prediction using both grid and average methods105# fig, axes = plt.subplots(2, 2, figsize=(15, 15))106 107# # Grid method visualization108# original_grid = combine_slices(image[0, ..., 0], method='grid')109# axes[0, 0].imshow(original_grid, cmap='gray')110# axes[0, 0].set_title('Original Image (Grid)')111# axes[0, 0].axis('off')112 113# prediction_grid = combine_slices(binary_prediction, method='grid')114# axes[0, 1].imshow(prediction_grid, cmap='gray')115# axes[0, 1].set_title('Predicted Mask (Grid)')116# axes[0, 1].axis('off')117 118# # Average method visualization119# original_avg = combine_slices(image[0, ..., 0], method='average')120# axes[1, 0].imshow(original_avg, cmap='gray')121# axes[1, 0].set_title('Original Image (Average)')122# axes[1, 0].axis('off')123 124# prediction_avg = combine_slices(binary_prediction, method='average')125# axes[1, 1].imshow(prediction_avg, cmap='gray')126# axes[1, 1].set_title('Predicted Mask (Average)')127# axes[1, 1].axis('off')128 129# plt.tight_layout()130 131# # Save the plot to a BytesIO object and encode it as base64132# img_buffer = BytesIO()133# plt.savefig(img_buffer, format='png', bbox_inches='tight', dpi=150)134# img_buffer.seek(0)135# img_str = base64.b64encode(img_buffer.getvalue()).decode()136# plt.close()137 138# # Return the image as a data URL139# return f"data:image/png;base64,{img_str}"140 141def predict_and_visualize(file_obj):142 # Save the uploaded file to the upload folder manually143 filename = secure_filename(file_obj.name)144 file_path = os.path.join(UPLOAD_FOLDER, filename)145 146 # Write the file content to disk147 with open(file_path, 'wb') as f:148 f.write(file_obj.read())149 150 # Copy the file to the static folder for display (optional)151 copy_to_static(filename)152 153 # Preprocess the image154 image = preprocess_image(file_path)155 prediction = model.predict(image)156 157 # Handle the model output158 bridge_output, logits = prediction159 160 # Use the logits (second output) for visualization161 binary_prediction = (logits[0, ..., 0] > 0.5).astype(np.uint8)162 163 # Save the prediction mask164 mask_filename = 'mask_' + filename165 mask_path = os.path.join(UPLOAD_FOLDER, mask_filename)166 np.save(mask_path, binary_prediction)167 copy_to_static(mask_filename) # Copy mask to static folder168 169 # Visualize original image and prediction using both grid and average methods170 fig, axes = plt.subplots(2, 2, figsize=(15, 15))171 172 # Grid method visualization173 original_grid = combine_slices(image[0, ..., 0], method='grid')174 axes[0, 0].imshow(original_grid, cmap='gray')175 axes[0, 0].set_title('Original Image (Grid)')176 axes[0, 0].axis('off')177 178 prediction_grid = combine_slices(binary_prediction, method='grid')179 axes[0, 1].imshow(prediction_grid, cmap='gray')180 axes[0, 1].set_title('Predicted Mask (Grid)')181 axes[0, 1].axis('off')182 183 # Average method visualization184 original_avg = combine_slices(image[0, ..., 0], method='average')185 axes[1, 0].imshow(original_avg, cmap='gray')186 axes[1, 0].set_title('Original Image (Average)')187 axes[1, 0].axis('off')188 189 prediction_avg = combine_slices(binary_prediction, method='average')190 axes[1, 1].imshow(prediction_avg, cmap='gray')191 axes[1, 1].set_title('Predicted Mask (Average)')192 axes[1, 1].axis('off')193 194 plt.tight_layout()195 196 # Save the plot to a BytesIO object and encode it as base64197 img_buffer = BytesIO()198 plt.savefig(img_buffer, format='png', bbox_inches='tight', dpi=150)199 img_buffer.seek(0)200 img_str = base64.b64encode(img_buffer.getvalue()).decode()201 plt.close()202 203 # Return the image as a data URL204 return f"data:image/png;base64,{img_str}"205 206 207def secure_filename(filename):208 """Sanitize the filename to prevent directory traversal attacks"""209 return os.path.basename(filename)210 211# Create the Gradio interface212iface = gr.Interface(213 fn=predict_and_visualize,214 inputs=gr.File(file_types=[".npy"], label="Upload Numpy Image (.npy)"),215 outputs=gr.Image(label="Prediction Results"),216 title="3D Image Prediction",217 description="Upload a 3D NumPy (.npy) image file to get predictions and visualizations."218)219 220# Launch the app221if __name__ == "__main__":222 iface.launch()223