import nibabel as nib
import matplotlib.pyplot as plt
import gradio as gr

def visualize_nii_gz(file_path):
    # Load the NIfTI file
    img = nib.load(file_path.name)
    data = img.get_fdata()
    
    # Get the middle slice of the data in all three dimensions
    x_slice = data[data.shape[0] // 2, :, :]
    y_slice = data[:, data.shape[1] // 2, :]
    z_slice = data[:, :, data.shape[2] // 2]
    
    # Plot the slices
    fig = plt.figure(figsize=(12, 4))
    
    ax1 = fig.add_subplot(131)
    ax1.imshow(x_slice.T, cmap="gray", origin="lower")
    ax1.set_title("X slice")
    
    ax2 = fig.add_subplot(132)
    ax2.imshow(y_slice.T, cmap="gray", origin="lower")
    ax2.set_title("Y slice")
    
    ax3 = fig.add_subplot(133)
    ax3.imshow(z_slice.T, cmap="gray", origin="lower")
    ax3.set_title("Z slice")
    
    # Return the figure
    plt.tight_layout()
    return fig

# Create the Gradio interface
interface = gr.Interface(
    fn=visualize_nii_gz, 
    inputs="file", 
    outputs="plot",
    live=True,
    title="3D NIfTI Visualizer",
    description="Upload a .nii.gz file to visualize its slices in 3D."
)

interface.launch()