Create app.py

app.py

@@ -0,0 +1,208 @@
+import streamlit as st
+import torch
+import os
+import tempfile
+import time
+
+# nnU-Net and visualization imports
+from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+import pyvista as pv
+import nibabel as nib
+import numpy as np
+from matplotlib import cm
+from matplotlib.colors import ListedColormap
+from stpyvista import stpyvista
+
+# --- Caching the nnU-Net Predictor ---
+# This is crucial for performance. The model is loaded once and stored in memory.
+@st.cache_resource
+def load_predictor(model_folder):
+    """
+    Loads and initializes the nnUNetPredictor.
+    The @st.cache_resource decorator ensures this function is only run once.
+    """
+    st.write("Initializing nnU-Net predictor... (This may take a moment)")
+    
+    # Instantiate the predictor
+    predictor = nnUNetPredictor(
+        tile_step_size=0.5,
+        use_gaussian=True,
+        use_mirroring=True,
+        perform_everything_on_device=True,
+        device=torch.device('cuda', 0) if torch.cuda.is_available() else torch.device('cpu'),
+        verbose=False,
+        verbose_preprocessing=False,
+        allow_tqdm=True
+    )
+    
+    # Initialize from the trained model folder
+    try:
+        predictor.initialize_from_trained_model_folder(
+            model_folder,
+            use_folds=(0,),  # Assuming you want to use fold 0
+            checkpoint_name='checkpoint_final.pth',
+        )
+        st.success("nnU-Net predictor initialized successfully!")
+        return predictor
+    except Exception as e:
+        st.error(f"Failed to initialize predictor from {model_folder}. Error: {e}")
+        return None
+
+# --- Visualization Function (from your script) ---
+def generate_visualization(base_image_path, mask_path):
+    """
+    Generates a PyVista plot of the base image and the segmentation mask.
+    """
+    # Load base CT scan
+    img = nib.load(base_image_path)
+    img_data = img.get_fdata()
+    img_data = (img_data - np.min(img_data)) / np.ptp(img_data)  # Normalize 0–1
+
+    # Load segmentation mask
+    mask = nib.load(mask_path)
+    mask_data = mask.get_fdata().astype(np.uint8)
+
+    # Label dictionary (from your script)
+    label_dict = {
+        1: "Lower Jawbone", 2: "Upper Jawbone", 3: "Left Inferior Alveolar Canal",
+        4: "Right Inferior Alveolar Canal", 5: "Left Maxillary Sinus", 6: "Right Maxillary Sinus",
+        7: "Pharynx", 8: "Bridge", 9: "Crown", 10: "Implant", 11: "Upper Right Central Incisor",
+        12: "Upper Right Lateral Incisor", 13: "Upper Right Canine", 14: "Upper Right First Premolar",
+        15: "Upper Right Second Premolar", 16: "Upper Right First Molar", 17: "Upper Right Second Molar",
+        18: "Upper Right Third Molar", 21: "Upper Left Central Incisor",
+        22: "Upper Left Lateral Incisor", 23: "Upper Left Canine", 24: "Upper Left First Premolar",
+        25: "Upper Left Second Premolar", 26: "Upper Left First Molar", 27: "Upper Left Second Molar",
+        28: "Upper Left Third Molar", 31: "Lower Left Central Incisor",
+        32: "Lower Left Lateral Incisor", 33: "Lower Left Canine", 34: "Lower Left First Premolar",
+        35: "Lower Left Second Premolar", 36: "Lower Left First Molar", 37: "Lower Left Second Molar",
+        38: "Lower Left Third Molar", 41: "Lower Right Central Incisor",
+        42: "Lower Right Lateral Incisor", 43: "Lower Right Canine", 44: "Lower Right First Premolar",
+        45: "Lower Right Second Premolar", 46: "Lower Right First Molar", 47: "Lower Right Second Molar",
+        48: "Lower Right Third Molar"
+    }
+    
+    # Generate color map
+    num_labels = max(label_dict.keys()) + 1
+    colors = np.vstack([
+        [[0, 0, 0, 0]],
+        cm.get_cmap('tab20b')(np.linspace(0, 1, 20)),
+        cm.get_cmap('tab20c')(np.linspace(0, 1, 20)),
+        cm.get_cmap('gist_rainbow')(np.linspace(0, 1, num_labels))
+    ])[:, :4]
+    colors = colors[:num_labels]
+    colormap = ListedColormap(colors)
+
+    # Wrap data in PyVista objects
+    vol_img = pv.wrap(img_data)
+    vol_mask = pv.wrap(mask_data)
+    
+    # Create plotter
+    plotter = pv.Plotter(window_size=[800, 800])
+    plotter.add_volume(vol_img, cmap="bone", opacity="sigmoid", name="CT Scan")
+    plotter.add_volume(
+        vol_mask,
+        cmap=colormap,
+        opacity=[0, 0.5], # Make label 0 transparent
+        mapper='gpu', # Use GPU for better performance
+        name="Segmentation Mask"
+    )
+    plotter.camera_position = 'xy'
+    
+    return plotter
+
+
+# --- Main Streamlit App ---
+def main():
+    st.set_page_config(layout="wide", page_title="nnU-Net Inference App")
+
+    st.title("🦷 nnU-Net Inference and 3D Visualization")
+    st.markdown("Upload a medical image, run nnU-Net for segmentation, and visualize the results in 3D.")
+
+    # --- Sidebar for Inputs ---
+    st.sidebar.header("1. Configure Model")
+    # IMPORTANT: Update this path to your default nnU-Net results folder
+    default_model_path = "/path/to/your/nnUNet_results/Dataset114_ToothFairy2/nnUNetTrainer__nnUNetPlans__3d_fullres"
+    model_folder = st.sidebar.text_input(
+        "Enter path to trained model folder:",
+        value=default_model_path
+    )
+
+    if not os.path.isdir(model_folder):
+        st.sidebar.error("Model folder not found. Please provide a valid path.")
+        st.stop()
+
+    # Load the model (will be cached)
+    predictor = load_predictor(model_folder)
+    if predictor is None:
+        st.stop()
+
+    st.sidebar.header("2. Upload Image")
+    uploaded_file = st.sidebar.file_uploader(
+        "Choose a NIfTI file (.nii.gz)",
+        type=['nii.gz']
+    )
+
+    # --- Main Panel for Execution and Visualization ---
+    if uploaded_file is not None:
+        if st.sidebar.button("✨ Run Prediction and Visualize"):
+            # Use a temporary directory for safety and automatic cleanup
+            with tempfile.TemporaryDirectory() as temp_dir:
+                input_dir = os.path.join(temp_dir, 'input')
+                output_dir = os.path.join(temp_dir, 'output')
+                os.makedirs(input_dir, exist_ok=True)
+                os.makedirs(output_dir, exist_ok=True)
+
+                # Save the uploaded file to the temp input directory
+                # The filename needs the _0000 suffix for nnU-Net's default file prediction
+                base_name = uploaded_file.name.replace(".nii.gz", "")
+                input_file_path = os.path.join(input_dir, f"{base_name}_0000.nii.gz")
+                
+                with open(input_file_path, "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                
+                st.info(f"File '{uploaded_file.name}' saved to temporary location.")
+
+                # --- Run Prediction ---
+                with st.spinner("🧠 Running nnU-Net inference... This can take a while."):
+                    start_time = time.time()
+                    
+                    # We use predict_from_files as it's the most efficient for file-based workflows
+                    predictor.predict_from_files(
+                        input_dir,
+                        output_dir,
+                        save_probabilities=False,
+                        overwrite=True,
+                        num_processes_preprocessing=2,
+                        num_processes_segmentation_export=2
+                    )
+                    
+                    end_time = time.time()
+                    st.success(f"Inference complete! 🎉 (Time taken: {end_time - start_time:.2f} seconds)")
+
+                # Find the output file
+                output_files = os.listdir(output_dir)
+                if not output_files:
+                    st.error("Prediction failed. No output file was generated.")
+                    st.stop()
+                
+                output_mask_path = os.path.join(output_dir, output_files[0])
+                
+                # --- Generate Visualization ---
+                with st.spinner("🎨 Generating 3D visualization..."):
+                    plotter = generate_visualization(input_file_path, output_mask_path)
+                    stpyvista(plotter, key="pv_plot")
+
+                # --- Provide Download Link for the Mask ---
+                with open(output_mask_path, "rb") as f:
+                    st.download_button(
+                        label="⬇️ Download Segmentation Mask",
+                        data=f,
+                        file_name=f"predicted_{uploaded_file.name}",
+                        mime="application/gzip"
+                    )
+
+    else:
+        st.info("Please upload a file to begin.")
+
+if __name__ == '__main__':
+    main()