Spaces:
Running
Running
File size: 3,539 Bytes
c6d85c0 |
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 |
# Import needed libraries
import streamlit as st
import torchvision
from medigan import Generators
from torchvision.transforms.functional import to_pil_image
from torchvision.utils import make_grid
# Define the GAN models available in the app
model_ids = [
"00001_DCGAN_MMG_CALC_ROI",
"00002_DCGAN_MMG_MASS_ROI",
"00003_CYCLEGAN_MMG_DENSITY_FULL",
"00004_PIX2PIX_MMG_MASSES_W_MASKS",
"00019_PGGAN_CHEST_XRAY" # New model
]
def main():
st.title("MEDIGAN Medical Image Data Generator")
# Add dropdown widget for model selection to the sidebar
model_id = st.sidebar.selectbox("Select Model ID", model_ids)
# Add number image selector to the sidebar
num_images = st.sidebar.number_input(
"Number of Images", min_value=1, max_value=7, value=1, step=1
)
# Add generate button to the sidebar
if st.sidebar.button("Generate Images"):
generate_images(num_images, model_id)
# Task 5.4.9: Copy the torch_images function from this notebook to app.py.
def torch_images(num_images, model_id):
generators = Generators()
dataloader = generators.get_as_torch_dataloader(
model_id=model_id,
install_dependencies=True,
num_samples=num_images,
prefetch_factor=None,
)
images = []
for batch_idx, data_dict in enumerate(dataloader):
image_list = []
for i in data_dict:
if "sample" in i:
sample = data_dict.get("sample")
if sample.dim() == 4:
sample = sample.squeeze(0).permute(2, 0, 1)
sample = to_pil_image(sample).convert("RGB")
# Convert the image to a PyTorch tensor
transform = torchvision.transforms.Compose(
[
torchvision.transforms.ToTensor(),
]
)
# Apply the transform to your PIL image
sample = transform(sample)
image_list.append(sample)
# Preprocess the mask
if "mask" in i:
mask = data_dict.get("mask")
if mask.dim() == 4:
mask = mask.squeeze(0).permute(2, 0, 1)
mask = to_pil_image(mask).convert("RGB")
mask = transform(mask)
image_list.append(mask)
# Organize the grid to have 'sample' images per row
Grid = make_grid(image_list, nrow=2)
# Change Grid tensor to be a consistent shape
# The Grid tensor has shape [1, 128, 128, 1] in some models
if Grid.dim() == 4:
# Remove the singleton batch dimension
Grid = Grid.squeeze(0)
if Grid.size(-1) == 1:
# Remove the singleton channel dimension (assuming grayscale)
Grid = Grid.squeeze(-1)
else:
raise ValueError("Expected a single channel (grayscale) image.")
# Convert the tensor grid to a PIL Image for display
img = torchvision.transforms.ToPILImage()(Grid)
images.append(img)
return images
def generate_images(num_images, model_id):
st.subheader("Generated Images:")
images = torch_images(num_images, model_id) # Pass the correct parameters
for i in range(len(images)):
st.image(
images[i],
caption=f"Generated Image {i+1} (Model ID: {model_id})",
use_container_width=True,
)
if __name__ == "__main__":
main() # Task 5.4.4: Call the main function
|