Spaces:
Sleeping
Sleeping
| import os | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import plotly.graph_objects as go | |
| from mpl_toolkits.mplot3d import Axes3D | |
| from skimage import measure | |
| from mpl_toolkits.mplot3d.art3d import Poly3DCollection | |
| from lungmask import LMInferer | |
| import SimpleITK as sitk | |
| import gradio as gr | |
| # --- Lung Segmentation Functions --- | |
| def process_dcm_file(file_path, inferer): | |
| """Loads a DCM file, performs lung segmentation, and returns the mask.""" | |
| input_image = sitk.ReadImage(file_path) | |
| segmentation = inferer.apply(input_image) | |
| newseg = segmentation.reshape(512, 512) # Assuming 512x512 images | |
| return newseg | |
| def segment_lungs_from_dicom(dcm_folder): | |
| """Segments lungs from DICOM files in a folder and returns a 3D volume.""" | |
| if not os.path.exists(dcm_folder) or not os.path.isdir(dcm_folder): | |
| raise ValueError("Invalid DICOM folder path.") | |
| inferer = LMInferer() | |
| segmentation_masks = [] | |
| for filename in os.listdir(dcm_folder): | |
| if filename.endswith(".dcm"): | |
| file_path = os.path.join(dcm_folder, filename) | |
| mask = process_dcm_file(file_path, inferer) | |
| segmentation_masks.append(mask) | |
| volume = np.stack(segmentation_masks, axis=0) | |
| return volume | |
| # --- 3D Visualization Function --- | |
| def plot_3d_lungs(lungs_volume, threshold=0.5): | |
| """Creates an interactive 3D plot of segmented lungs using Plotly (upright).""" | |
| verts, faces, normals, values = measure.marching_cubes(lungs_volume.transpose(2, 1, 0), threshold) | |
| # Apply rotation to make lungs upright | |
| # Assuming you want to rotate 90 degrees counter-clockwise around the X-axis | |
| rotation_angle_degrees = -90 | |
| rotation_angle_radians = np.radians(rotation_angle_degrees) | |
| rotation_matrix = np.array([[1, 0, 0], | |
| [0, np.cos(rotation_angle_radians), -np.sin(rotation_angle_radians)], | |
| [0, np.sin(rotation_angle_radians), np.cos(rotation_angle_radians)]]) | |
| rotated_verts = np.dot(verts, rotation_matrix) | |
| x, y, z = zip(*rotated_verts) # Use rotated vertices | |
| i, j, k = zip(*faces) | |
| mesh = go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, opacity=0.7, color='lightblue') | |
| fig = go.Figure(data=[mesh]) | |
| fig.update_layout(scene_aspectmode='data') # Maintain aspect ratio | |
| return fig | |
| # --- Gradio Interface --- | |
| def process_and_visualize(selected_folder): | |
| if selected_folder not in ["tumor", "lung", "tumor2"]: | |
| return "Invalid folder selection." # Handle invalid input | |
| volume = segment_lungs_from_dicom(selected_folder) | |
| visualization = plot_3d_lungs(volume) | |
| return visualization | |
| inputs = gr.Dropdown(choices=["tumor", "lung", "tumor2"], label="Select DICOM Folder") | |
| output = gr.Plot(label="3D Segmented Lungs") | |
| iface = gr.Interface( | |
| fn=process_and_visualize, | |
| inputs=inputs, | |
| outputs=output, | |
| title="3D Lung Segmentation Visualization", | |
| description="Visualize segmented lungs from DICOM images.", | |
| ) | |
| iface.launch() |