VersatileSegmentation/create_3sets_tree_repr_partial.py at main · hzaoheng123/VersatileSegmentation · 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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Aug 18 13:27:54 2023

@author: xychen
"""

import os
import pickle
import numpy as np
import SimpleITK as sitk

datasets_for_training = ["amos", "btcv", "flare22"]
modalities_for_training = ["ct", "t1w"]

data_path = "/mnt/yapdata/data/xychen/awesome/MM/labelsTr"
# data_path = "/mnt/8TBSSD/chexiao/awesome/MM/labelsTr"
files = os.listdir(data_path)

hierarchy = {}
for file in files:

    label = sitk.ReadImage(os.path.join(data_path, file))
    label = sitk.GetArrayFromImage(label)

    categories_present = np.sort(np.unique(np.round(label).astype(np.int32)))

    print("Current file = ", file, " have annotations ", categories_present)

    modality = file.split('_')[1]
    dataset = file.split('_')[0]

    if dataset not in datasets_for_training or modality not in modalities_for_training:
        continue

    for category in categories_present:

        if category not in hierarchy.keys():
            hierarchy[category] = {}

        if modality not in hierarchy[category].keys():
            hierarchy[category][modality] = {}

        if dataset not in hierarchy[category][modality].keys():
            hierarchy[category][modality][dataset] = []

        if file not in hierarchy[category][modality][dataset]:
            hierarchy[category][modality][dataset].append(file)


with open('tree_repr_3sets_mm_partial.pkl', 'wb') as f:
    pickle.dump(hierarchy, f)