Spaces:
Sleeping
Sleeping
File size: 2,446 Bytes
1f2b3fa |
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 |
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")
|