Spaces:
Sleeping
Sleeping
File size: 3,655 Bytes
e9406c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import os
import uuid
import subprocess
import re
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def get_least_used_gpu(default_gpu=None):
if default_gpu is None:
try:
available_gpus_str = os.getenv("AVAILABLE_GPUS", "")
available_gpus = [int(x) for x in available_gpus_str.split(",") if x.strip().isdigit()]
if not available_gpus:
raise ValueError("No available GPUs specified.")
result = subprocess.check_output(
["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
universal_newlines=True
)
mem_usages = [int(x) for x in result.strip().split("\n")]
least_used_gpu = min(available_gpus, key=lambda i: mem_usages[i])
return str(least_used_gpu)
except Exception as e:
print("⚠️ Failed to get GPU info, defaulting to 0:", e)
return "0"
else:
return str(default_gpu)
def run_auto_segmentation(input_path, session_dir, model):
"""
Run auto segmentation model using Apptainer inside the given session directory.
"""
subfolder_name = "ct"
input_case_dir = os.path.join(session_dir, "inputs")
outputs_root = os.path.join(session_dir, "outputs")
input_case_ct_dir = os.path.join(input_case_dir, subfolder_name)
os.makedirs(input_case_ct_dir, exist_ok=True)
os.makedirs(outputs_root, exist_ok=True)
input_filename = os.path.basename(input_path)
container_input_path = os.path.join(input_case_ct_dir, input_filename)
os.system(f"cp {input_path} {container_input_path}")
conda_activate_cmd = ""
conda_path = os.getenv("CONDA_ACTIVATE_PATH", "/opt/anaconda3/etc/profile.d/conda.sh")
epai_env_name = os.getenv("CONDA_ENV_EPAI", "epai")
suprem_sandbox_path = os.getenv("SUPREM_SANDBOX_PATH", "")
epai_script_path = os.getenv("EPAI_SCRIPT_PATH", "")
if model == 'SuPreM':
container_path = suprem_sandbox_path
print(input_case_dir, outputs_root)
apptainer_cmd = [
"apptainer", "run", "--nv",
"-B", f"{input_case_dir}:/workspace/inputs",
"-B", f"{outputs_root}:/workspace/outputs",
container_path
]
elif model == 'ePAI':
conda_activate_cmd = f"source {conda_path} && conda activate {epai_env_name} &&"
apptainer_cmd = ["bash", epai_script_path, session_dir]
else:
print(f"[ERROR] Unknown model: {model}")
return None
selected_gpu = get_least_used_gpu()
apptainer_cmd = ["CUDA_VISIBLE_DEVICES=" + selected_gpu] + apptainer_cmd
print(apptainer_cmd)
try:
print(f"[INFO] Running {model} auto segmentation for file: {input_filename}")
full_cmd = f"{conda_activate_cmd} {' '.join(apptainer_cmd)}"
subprocess.run(full_cmd, shell=True, executable="/bin/bash", check=True)
except subprocess.CalledProcessError as e:
print(f"[ERROR] {model} inference failed:", e)
return None
if model == 'SuPreM':
output_path = os.path.join(outputs_root, subfolder_name, "segmentations")
if not os.path.exists(output_path):
print("[ERROR] Output mask not found at:", output_path)
return None
elif model == 'ePAI':
output_path = os.path.join(outputs_root, subfolder_name, "combined_labels.nii.gz")
if not os.path.exists(output_path):
print("[ERROR] Output mask not found at:", output_path)
return None
output_path = os.path.join(outputs_root, subfolder_name)
return output_path
|