Benja24/TotalSegmentatorApp

1import os2import shutil3import subprocess4import uuid5import zipfile6import nibabel as nib7import gradio as gr8def run_segmentation(9    uploaded_file_path, selected_task, device, fast, roi_subset,10    robust_crop, preview, ml, statistics, radiomics, progress=gr.Progress(track_tqdm=True)11):12    import nibabel as nib13    import numpy as np14 15    job_id = str(uuid.uuid4())16    input_filename = f"input_{job_id}.nii.gz"17    output_folder = f"segmentations_{job_id}"18 19    shutil.copy(uploaded_file_path.name, input_filename)20 21    command = ["TotalSegmentator", "-i", input_filename, "-o", output_folder]22    if selected_task: command.extend(["--task", selected_task])23    if device: command.extend(["--device", device])24    if fast: command.append("--fast")25    if roi_subset: command.extend(["--roi_subset"] + roi_subset.strip().split())26    if robust_crop: command.append("--robust_crop")27    if preview: command.append("--preview")28    if ml: command.append("--ml")29    if statistics: command.append("--statistics")30    if radiomics:31        try:32            import radiomics33            command.append("--radiomics")34        except ImportError:35            print("Radiomics no disponible")36 37    try:38        subprocess.run(command, check=True)39    except subprocess.CalledProcessError:40        return None, None, 0, "Error durante segmentación.", gr.update(visible=False), None41 42    # Filtrar solo segmentaciones que no estén vacías43    seg_files = []44    for f in os.listdir(output_folder):45        if not f.endswith('.nii.gz'):46            continue47        full_path = os.path.join(output_folder, f)48        try:49            data = nib.load(full_path).get_fdata()50            if np.any(data):  # Tiene contenido51                seg_files.append(f)52        except Exception as e:53            print(f"Error leyendo {f}: {e}")54 55    if not seg_files:56        return (57            None, None, 0,58            "Segmentación completada, pero todas las máscaras están vacías.",59            gr.update(choices=[], visible=False),60            output_folder61        )62 63    # Crear ZIP solo con segmentaciones no vacías64    zip_filename = f"segmentations_{job_id}.zip"65    with zipfile.ZipFile(zip_filename, "w") as zipf:66        for f in seg_files:67            full_path = os.path.join(output_folder, f)68            zipf.write(full_path, f)69 70    first_seg_path = os.path.join(output_folder, seg_files[0])71    seg_data = nib.load(first_seg_path).get_fdata()72    depth = seg_data.shape[2]73 74    return (75        zip_filename,76        None, None, None,  # Placeholders for the three views (will be updated)77        gr.update(maximum=seg_data.shape[2]-1, value=seg_data.shape[2]//2),  # axial slider78        gr.update(maximum=seg_data.shape[1]-1, value=seg_data.shape[1]//2),  # coronal slider79        gr.update(maximum=seg_data.shape[0]-1, value=seg_data.shape[0]//2),  # sagittal slider80        "Segmentación completada.",81        gr.update(choices=seg_files, value=seg_files[0], visible=True),82        output_folder83    )