preprocessing_tractogram.py

# /auto/home/users/v/l/vlahy/code/preprocessing_tractogram.py

import sys
from pathlib import Path
import elikopy

# write skipped when missing files or tracking crash
def write_skipped(out_dir: Path, reason: str):
    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / "SKIPPED.txt", "w") as f:
        f.write(reason + "\n")
    print(f"[SKIPPED] {out_dir.name} → {reason}")

def main():
    if len(sys.argv) != 4:
        print("Usage: python preprocessing_tractogram.py <root_in> <root_out> <subject_id>")
        sys.exit(0)  

    root_in = Path(sys.argv[1])
    root_out = Path(sys.argv[2])
    sub = sys.argv[3]

    out_sub = root_out / sub
    out_sub.mkdir(parents=True, exist_ok=True)

    subj_dir = root_in / "subjects" / sub

    # ----------- Files to check -----------

    raw_dir = subj_dir / "dMRI" / "raw"
    odf_dir = subj_dir / "dMRI" / "ODF" / "MSMT-CSD"
    mask_dir = subj_dir / "masks"

    raw_nii  = raw_dir / f"{sub}_raw_dmri.nii.gz"
    raw_bval = raw_dir / f"{sub}_raw_dmri.bval"
    raw_bvec = raw_dir / f"{sub}_raw_dmri.bvec"

    wm_odf = odf_dir / f"{sub}_MSMT-CSD_WM_ODF.nii.gz"
    brain_mask = mask_dir / f"{sub}_brain_mask.nii.gz"

    # ----------- Checks -----------

    missing = []

    for f in [raw_nii, raw_bval, raw_bvec]:
        if not f.exists():
            missing.append(str(f))

    if not wm_odf.exists():
        missing.append(str(wm_odf))

    if not brain_mask.exists():
        missing.append(str(brain_mask))

    if missing:
        reason = "Missing required files:\n" + "\n".join(missing)
        write_skipped(out_sub, reason)
        sys.exit(0)  # VERY IMPORTANT: no FAIL

    # ----------- Tracking -----------

    try:
        study = elikopy.core.Elikopy(str(root_in), slurm=False, cuda=False)

        study.tracking(
            folder_path=str(root_in),
            save_as_trk=True,
            output_dir=str(out_sub) + "/",
            streamline_number=1000000,
            patient_list_m=[sub],
        )

        with open(out_sub / "DONE.txt", "w") as f:
            f.write("Tracking completed successfully\n")

        print(f"[DONE] {sub}")

    except Exception as e:
        write_skipped(out_sub, f"Tracking crashed:\n{str(e)}")
        sys.exit(0)  

if __name__ == "__main__":
    main()