Update app.py

app.py

@@ -11,92 +11,79 @@ model_ids = [
     "00002_DCGAN_MMG_MASS_ROI",
     "00003_CYCLEGAN_MMG_DENSITY_FULL",
     "00004_PIX2PIX_MMG_MASSES_W_MASKS",
-    "00019_PGGAN_CHEST_XRAY" # New model
+    "00019_PGGAN_CHEST_XRAY"
 ]
 
 def main():
-    st.title("MEDIGAN Medical Image Data Generator")
+    # Setup page configuration
+    st.set_page_config(page_title="MEDIGAN Generator", layout="wide")
+    
+    # Main page title and description
+    st.title("🧠 MEDIGAN Medical Image Generator")
+    st.markdown("""
+    **Generate synthetic medical images using GAN models.**  
+    🔍 Select model and parameters in the sidebar → Click **Generate Images**
+    """)
+    
+    # Sidebar controls
+    with st.sidebar:
+        st.header("⚙️ Settings")
+        model_id = st.selectbox("Select GAN Model", model_ids)
+        num_images = st.number_input("Number of Images", 1, 7, 1)
+        generate_btn = st.button("✨ Generate Images")
+    
+    # Main content area
+    if generate_btn:
+        with st.spinner(f"Generating {num_images} image(s) using {model_id}..."):
+            generate_images(num_images, model_id)
 
-    # 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,
+        num_workers=0,
+        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)
+    for _, data_dict in enumerate(dataloader):
+        batch_images = []
+        for tensor in data_dict.values():
+            if tensor.dim() == 4:
+                tensor = tensor.squeeze(0).permute(2, 0, 1)
+            img = to_pil_image(tensor).convert("RGB")
+            batch_images.append(img)
+        
+        # Create image grid
+        grid_tensor = make_grid(
+            [torchvision.transforms.ToTensor()(img) for img in batch_images],
+            nrow=2 if len(batch_images) > 1 else 1
+        )
+        grid_img = to_pil_image(grid_tensor)
+        images.append(grid_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,
-        )
+    # Clear previous results
+    st.empty()
+    
+    # Generate and display new images
+    images = torch_images(num_images, model_id)
+    
+    # Create columns for responsive layout
+    cols = st.columns(len(images))
+    
+    for col, img in zip(cols, images):
+        with col:
+            st.image(
+                img,
+                caption=f"Generated by: {model_id}",
+                width=300
+            )
+            st.markdown("---")
 
 if __name__ == "__main__":
-    main()  # Task 5.4.4: Call the main function
+    main()