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()