File size: 4,611 Bytes
c3e587d 0e28d84 c3e587d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | import streamlit as st
import numpy as np
import cv2
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing.image import img_to_array
from PIL import Image
import io
import zipfile
# Load multiple trained U-Net models
model_1 = tf.keras.models.load_model("pages/dyskera_unet_model.keras")
model_2 = tf.keras.models.load_model("pages/koilo_unet_model.keras")
model_3 = tf.keras.models.load_model("pages/meta_unet_model.keras")
model_4 = tf.keras.models.load_model("pages/para_unet_model.keras")
model_5 = tf.keras.models.load_model("pages/SI_unet_model.keras")
# Add more models as needed
# Function to preprocess the uploaded image
def preprocess_image(uploaded_image, target_size=(128, 128)):
image = Image.open(uploaded_image).convert('L') # Convert to grayscale
image = np.array(image) # Convert to numpy array
image = cv2.resize(image, target_size) # Resize the image
image = image.astype('float32') / 255.0 # Normalize
image = np.expand_dims(image, axis=-1) # Add channel dimension
image = np.expand_dims(image, axis=0) # Add batch dimension
return image
# Function to predict masks using the selected model
def predict_masks(image, model):
predictions = model.predict(image)
return predictions[0]
# Function to save masks as downloadable images
def save_image(mask, filename):
mask = (mask * 255).astype(np.uint8)
im = Image.fromarray(mask)
buf = io.BytesIO()
im.save(buf, format="PNG")
buf.seek(0)
return buf
# Streamlit Interface
st.title("U-Net Cell Segmentation with Multiple Models")
st.sidebar.info("Feel free to select other models from the pages above 🙂")
st.sidebar.write("""
### Instructions 🔧
Please upload Cell Image and choose appropriate U-net Model from Drop-down menu to get Predicted Cytoplasm Mask and Nuclei Mask. After getting and downloading Predicted Cytoplasm and Nuclei Mask image, head over to CNN or SVM model to get accurate prediction of Cancer Cell Type.
""")
# Upload an image
uploaded_image = st.file_uploader("Choose an image...", type=["bmp", "png", "jpg", "jpeg"])
# Model selection dropdown
model_option = st.selectbox(
"Select a Model:",
("U-Net Dyskeratotic", "U-Net Koilocyctotic", "U-Net Metaplastic", "U-Net Parabasal", "U-Net Superficial-Intermediate")
)
# Choose the corresponding model based on selection
if model_option == "U-Net Dyskeratotic":
selected_model = model_1
elif model_option == "U-Net Koilocyctotic":
selected_model = model_2
elif model_option == "U-Net Metaplastic":
selected_model = model_3
elif model_option == "U-Net Parabasal":
selected_model = model_4
elif model_option == "U-Net Superficial-Intermediate":
selected_model = model_5
if uploaded_image is not None:
st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
# Preprocess the image
image = preprocess_image(uploaded_image)
# Predict the masks using the selected model
predictions = predict_masks(image, selected_model)
# Display results using Matplotlib
fig, ax = plt.subplots(1, 3, figsize=(12, 8))
ax[0].imshow(image.squeeze(), cmap='gray')
ax[0].set_title('Original Image')
ax[1].imshow(predictions[:, :, 0], cmap='gray')
ax[1].set_title('Predicted Cytoplasm Mask')
ax[2].imshow(predictions[:, :, 1], cmap='gray')
ax[2].set_title('Predicted Nuclei Mask')
st.pyplot(fig)
# Convert the predicted masks to images for display
cytoplasm_mask_img = (predictions[:, :, 0] * 255).astype(np.uint8)
nuclei_mask_img = (predictions[:, :, 1] * 255).astype(np.uint8)
st.image(cytoplasm_mask_img, caption="Predicted Cytoplasm Mask", use_column_width=True)
st.image(nuclei_mask_img, caption="Predicted Nuclei Mask", use_column_width=True)
# Save the predicted masks for download
cytoplasm_mask_download = save_image(predictions[:, :, 0], "predicted_cytoplasm_mask.png")
nuclei_mask_download = save_image(predictions[:, :, 1], "predicted_nuclei_mask.png")
# Combine the two mask files into a zip
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w") as zip_file:
zip_file.writestr("predicted_cytoplasm_mask.png", cytoplasm_mask_download.getvalue())
zip_file.writestr("predicted_nuclei_mask.png", nuclei_mask_download.getvalue())
zip_buffer.seek(0)
# Provide a download button for the zipped masks
st.download_button(
label="Download Predicted Cytoplasm and Nuclei Masks",
data=zip_buffer,
file_name="predicted_masks.zip",
mime="application/zip"
) |