LargeBrainModel/dataLoaderLazy.py at master · GabrieleBan/LargeBrainModel · GitHub

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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import json
import platform
import nibabel as nib

def get_root_path():

    if platform.system() == "Windows":
        return Path(r"data03\cartellaGB\quasiRaw\quasiRawFomo60k\quasi-raw")
    else:
        return Path("/data03/cartellaGB/quasiRaw/quasiRawFomo60k/quasi-raw")

def niiLoader(nii_path):
    if isinstance(nii_path, str):

        nii_path = nii_path.replace("\\", "/")

    nii_path = Path(nii_path).resolve()

    if not nii_path.exists():
        print(f"❌ File non trovato: {nii_path}")
        return None

    img = nib.load(str(nii_path))
    return {
        "data": img.get_fdata(),
        "affine": img.affine
    }
from pathlib import Path
import re

def dataLoaderLazy(root_dir, output_file, pattern="*preproc-quasiraw_*.nii.gz", maxImg=float('inf')):
    results = {}
    count = 0
    scratch_base = Path.home()
    print(">>> Inizio dataLoader")
    print("Working dir:", Path.cwd())
    print("Scanning root:", root_dir)

    # forza lista per vedere se è vuota
    nii_files = list(root_dir.rglob("*preproc-quasiraw_*.nii.gz"))
    print(f"Found {len(nii_files)} matching files.")

    if not nii_files:
        print("⚠️ Nessun file trovato! Verifica il path e il pattern.")
    nii_files = sorted(root_dir.rglob(pattern))

    for nii_path in nii_files:
        if count >= maxImg:
            break

        fname = nii_path.name

        match = re.match(
            r"(?P<sub>sub-[^_]+)_?(?P<ses>ses-[^_]+)?_?(?P<run>run-[^_]+)?_?preproc-quasiraw_(?P<modality>[^_]+)(?P<mask>_mask)?\.nii\.gz",
            fname
        )

        if not match:
            continue

        sub = match.group("sub")
        ses = match.group("ses") or "ses-1"
        run = match.group("run") or "run-1"
        modality = match.group("modality")
        is_mask = match.group("mask") is not None

        # Inizializza soggetto se nuovo
        if sub not in results:
            results[sub] = {
                "subject": sub,
                "data": []
            }

        subj = results[sub]


        entry = next(
            (item for item in subj["data"] if item["session"] == ses and item["run"] == run and item["modality"] == modality),
            None
        )
        if entry is None:
            entry = {
                "session": ses,
                "run": run,
                "modality": modality,
                "image_path": None,
                "mask_path": None
            }
            subj["data"].append(entry)
        print("entry")
        rel_path = ("/"+str(nii_path.relative_to(scratch_base)))
        if is_mask:
            entry["mask_path"] = rel_path
        else:
            entry["image_path"] = rel_path
            count += 1


        next_sub = None
        next_idx = nii_files.index(nii_path) + 1
        if next_idx >= len(nii_files):
            should_save = True
        else:
            next_fname = nii_files[next_idx].name
            m = re.match(r"(?P<sub>sub-[^_]+)_", next_fname)
            next_sub = m.group("sub") if m else None
            should_save = next_sub != sub

        if should_save:
            print("saving")
            with open(output_file, "a", encoding="utf-8") as f:
                json.dump(subj, f)
                f.write("\n")
            del results[sub]
            print("done")

    return None

if __name__ == '__main__':
    root = Path.home() / "scratch" / "Fomo"
    output = Path("output.jsonl")
    dataLoaderLazy(root, output)