Spaces:
Sleeping
Sleeping
File size: 2,191 Bytes
a578142 b1969a7 a578142 |
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 |
from src.model.unet import UNet
import streamlit as st
import torch
from torchvision import transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2
from PIL import Image
import numpy as np
import config.configure as config
from src.pipelines.predict import predict_mask
model = UNet(3, 1, [64, 128, 256, 512])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.load_state_dict(torch.load(config.SAVE_MODEL_PATH, map_location=torch.device(device)))
# Set up transformations for the input image
transform = A.Compose([
A.Resize(224, 224, p=1.0),
ToTensorV2(),
])
# Streamlit app
def main():
st.title("MRI segmenation App")
# Upload image through Streamlit
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_image is not None:
# Display the uploaded and processed images side by side
col1, col2 = st.columns(2) # Using beta_columns for side-by-side layout
# Display the uploaded image in the first column
col1.header("Original Image")
col1.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
# Process the image (replace this with your processing logic)
processed_image = generate_image(uploaded_image)
# Display the processed image in the second column
col2.header("Processed Image")
col2.image(processed_image, caption="Segmented Image", use_column_width=True)
# Function to generate an image using the PyTorch model
def generate_image(uploaded_image):
# Load the uploaded image
input_image = Image.open(uploaded_image)
image = np.array(input_image).astype(np.float32) / 255.
# Apply transformations
input_tensor = transform(image=image)["image"].unsqueeze(0)
# Generate an image using the PyTorch model
mask = predict_mask(data=input_tensor, device=device, model=model, inference=True)
mask = mask[0].permute(1, 2, 0)
image = input_tensor[0].permute(1, 2, 0)
mask = image + mask*0.3
mask = mask.permute(2, 0, 1)
mask = transforms.ToPILImage()(mask)
return mask
if __name__ == "__main__":
main() |