Benja24/TotalSegmentatorApp

1import matplotlib.pyplot as plt 2import numpy as np3from matplotlib.widgets import Slider4from PIL import Image5import gradio as gr 6import subprocess7import os8import shutil9import uuid10import zipfile11import nibabel as nib12import pydicom13import uuid14 15def load_image(file):16    ext = os.path.splitext(file.name)[1].lower()17    if ext in [".gz", ".nii"]:18        nii = nib.load(file.name)19        data = nii.get_fdata()20    elif ext in [".dcm"]:21        ds = pydicom.dcmread(file.name)22        data = ds.pixel_array23        if data.ndim == 2:24            data = data[np.newaxis, :, :]  # agrega dimensión axial para manejar igual que volumen 3D25    else:26        raise ValueError("Formato no soportado. Usa .nii.gz o .dcm")27    return data28 29 30def plot_slice_np(data, plane, slice_idx):31    if plane == "axial":32        img = data[slice_idx, :, :]33    elif plane == "coronal":34        img = data[:, slice_idx, :]35    elif plane == "sagittal":36        img = data[:, :, slice_idx]37    else:38        raise ValueError("Plano inválido.")39 40    # Normalizar para que esté entre 0 y 255 uint841    img_norm = img - np.min(img)42    if np.max(img_norm) > 0:43        img_norm = img_norm / np.max(img_norm)44    img_uint8 = (img_norm * 255).astype(np.uint8)45 46    return np.rot90(img_uint8)47 48def process_and_visualize(image_file):49    volume = load_image(image_file)50    shape = volume.shape51    axial_center = shape[0] // 252    coronal_center = shape[1] // 253    sagittal_center = shape[2] // 254 55    return (56        gr.update(visible=True),57        volume,58        gr.update(minimum=0, maximum=shape[0]-1, value=axial_center),59        gr.update(minimum=0, maximum=shape[1]-1, value=coronal_center),60        gr.update(minimum=0, maximum=shape[2]-1, value=sagittal_center),61        plot_slice_np(volume, "axial", axial_center),62        plot_slice_np(volume, "coronal", coronal_center),63        plot_slice_np(volume, "sagittal", sagittal_center),64    )65 66def update_axial(volume, idx): return plot_slice_np(volume, "axial", idx)67def update_coronal(volume, idx): return plot_slice_np(volume, "coronal", idx)68def update_sagittal(volume, idx): return plot_slice_np(volume, "sagittal", idx)69 70def show_segmentation_slice(seg_path, slice_idx):71    seg_data = nib.load(seg_path).get_fdata()72    slice_img = seg_data[:, :, slice_idx]73    slice_norm = slice_img - np.min(slice_img)74    if np.max(slice_norm) > 0:75        slice_norm /= np.max(slice_norm)76    return (slice_norm * 255).astype(np.uint8)77def update_segmentation_view(seg_name, output_folder, volume, slice_idx, plane):78    import matplotlib.pyplot as plt79    import io80    from PIL import Image81    import numpy as np82 83    if seg_name is None or output_folder is None or volume is None:84        return None85 86    # Load segmentation87    seg_path = os.path.join(output_folder, seg_name)88    if not os.path.exists(seg_path):89        return None90    seg_data = nib.load(seg_path).get_fdata()91 92    # Get the appropriate slice93    if plane == "axial":94        img_slice = volume[:, :, slice_idx]95        mask_slice = seg_data[:, :, slice_idx]96    elif plane == "coronal":97        img_slice = volume[:, slice_idx, :]98        mask_slice = seg_data[:, slice_idx, :]99    elif plane == "sagittal":100        img_slice = volume[slice_idx, :, :]101        mask_slice = seg_data[slice_idx, :, :]102    else:103        return None104 105    # Normalize image106    img_norm = img_slice - np.min(img_slice)107    if np.max(img_norm) > 0:108        img_norm = img_norm / np.max(img_norm)109    base_img = (img_norm * 255).astype(np.uint8)110 111    # Create overlay112    fig, ax = plt.subplots(figsize=(5, 5))113    ax.imshow(base_img, cmap="gray")114    ax.imshow(mask_slice, cmap="jet", alpha=0.5)115    ax.axis("off")116 117    # Convert to PIL Image118    buf = io.BytesIO()119    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)120    plt.close(fig)121    buf.seek(0)122    overlay_img = Image.open(buf)123 124    return np.rot90(overlay_img)125 126def update_all_segmentation_views(seg_name, output_folder, volume):127    if seg_name is None or output_folder is None or volume is None:128        return [None]*6129 130    # Load segmentation to get dimensions131    seg_path = os.path.join(output_folder, seg_name)132    seg_data = nib.load(seg_path).get_fdata()133    134    # Calculate center slices for each plane135    axial_center = seg_data.shape[2] // 2136    coronal_center = seg_data.shape[1] // 2137    sagittal_center = seg_data.shape[0] // 2138    139    # Update all views140    axial_view = update_segmentation_view(seg_name, output_folder, volume, axial_center, "axial")141    coronal_view = update_segmentation_view(seg_name, output_folder, volume, coronal_center, "coronal")142    sagittal_view = update_segmentation_view(seg_name, output_folder, volume, sagittal_center, "sagittal")143    144    return [145        axial_view, coronal_view, sagittal_view,146        gr.update(maximum=seg_data.shape[2]-1, value=axial_center),147        gr.update(maximum=seg_data.shape[1]-1, value=coronal_center),148        gr.update(maximum=seg_data.shape[0]-1, value=sagittal_center)149    ]