Create app.py

app.py

@@ -0,0 +1,77 @@
+import streamlit as st
+from PIL import Image
+import torch
+from torchvision import transforms
+import numpy as np
+import os
+from osgeo import gdal
+
+# Load the pretrained model
+@st.cache(allow_output_mutation=True)
+def load_model():
+    model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
+                           pretrained=True, progress=True)
+    model.eval()
+    return model
+
+# Function to load large TIFF images
+def load_tiff_image(tiff_path):
+    try:
+        dataset = gdal.Open(tiff_path)
+        if dataset is None:
+            st.error("Failed to load the TIFF image. Please check the file format.")
+            return None
+        band = dataset.GetRasterBand(1)  # Assuming grayscale or single band
+        image = band.ReadAsArray()
+        return image
+    except Exception as e:
+        st.error(f"Error loading image: {e}")
+        return None
+
+# Preprocess image
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Resize((256, 256)),  # Resize image for model input
+        transforms.Normalize(mean=[0.485], std=[0.229])  # Normalize
+    ])
+    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+    return image_tensor
+
+# Post-process prediction to display
+def postprocess_prediction(pred):
+    pred = torch.sigmoid(pred)
+    pred = pred.squeeze().detach().numpy()  # Remove batch dimension
+    pred = (pred > 0.5).astype(np.uint8)  # Binary mask thresholding
+    return pred
+
+# Streamlit app
+st.title("TIFF Image Upload and Model Prediction")
+
+# Upload image
+uploaded_file = st.file_uploader("Upload a large TIFF image (up to 5GB)", type=["tiff"])
+
+if uploaded_file is not None:
+    with open("temp_image.tiff", "wb") as f:
+        f.write(uploaded_file.getbuffer())
+    
+    tiff_image = load_tiff_image("temp_image.tiff")
+    
+    if tiff_image is not None:
+        st.write("Original Image")
+        st.image(tiff_image, caption="Uploaded Image", use_column_width=True)
+        
+        model = load_model()
+        
+        image = Image.fromarray(tiff_image)
+        image_tensor = preprocess_image(image)
+        
+        with torch.no_grad():
+            prediction = model(image_tensor)
+        
+        pred_image = postprocess_prediction(prediction)
+        
+        st.write("Model Prediction")
+        st.image(pred_image, caption="Predicted Image", use_column_width=True)
+    
+    os.remove("temp_image.tiff")