ZAugNet/utils.py at main · VirtualEmbryo/ZAugNet · 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
125
126
127
128
129
130
import os
import shutil
import itertools
import tifffile as tiff

import numpy as np

import torch

EPS =  1e-12


def normalization(image0, image1, image2, cfg):
    if cfg.normalization == 'min_max':
        min_, max_ = cfg.min_max
        image0 = (image0 - min_)/(max_ - min_)
        image1 = (image1 - min_)/(max_ - min_)
        image2 = (image2 - min_)/(max_ - min_)


    elif cfg.normalization == 'zscore':
        image0 = (image0 - image0.float().mean()) / max(image0.float().std(), EPS)
        image1 = (image1 - image1.float().mean()) / max(image1.float().std(), EPS)
        image2 = (image2 - image2.float().mean()) / max(image2.float().std(), EPS)

    return image0, image1, image2


def create_zdataset(list_of_files:str, cfg:dict):
    """
    Generates a 'datasets' folder containing one dataset per channel. Each dataset comprises PyTorch tensors,
    with each tensor representing three consecutive 2D slices extracted from the input 5D image array.

    Args:
        list_of_files (str): File paths of the input images.
        cfg (dict): Configuration options including dataset size, model name, and validation split ratio.

    Returns:
        None
    """
    dataset_size = cfg.dataset_size
    model_name = cfg.model_name
    p_val = cfg.p_val

    triplets = np.empty((0, 5), dtype=int)

    root = '.'
    for i, file in enumerate(list_of_files):

        path = f"{root}/data/train/{file}"
        image = tiff.imread(path)

        format = 'y'
        len_shape = len(image.shape)
        if len_shape <= 2:
            raise ValueError('Error: images must be at least 3D.')
        elif len_shape == 3:
            image = image[None, :, None, :, :]
        elif len_shape == 4:
            if format == 'y':
                image = image[:, :, None, :, :]
            elif format == 'n':
                image = image[None, :, :, :, :]
            else:
                raise ValueError('Error: data file-format not recognized.')
        else:
            raise ValueError('Error: images must be at most 5D.')

        t, z = image.shape[0:2]

        if model_name == 'zaugnet+':
            all_triplets = np.array(list(itertools.combinations(range(z), 3)))
            selected_all_triplets = []
            for idx in range(len(all_triplets)):
                if (all_triplets[idx][2] - all_triplets[idx][0]) <= cfg.distance_triplets :
                    selected_all_triplets.append(all_triplets[idx])
            all_triplets = np.array(selected_all_triplets)

            for dt in range(t):
                rand_ = np.random.randint(0, len(all_triplets), dataset_size)

                selected_triplets = np.concatenate((np.repeat([[i,dt]], dataset_size, axis=0), all_triplets[rand_]), axis=1)
                triplets = np.concatenate((triplets, selected_triplets))


        else :
            for dt, dz in itertools.product(range(t), range(z - 2)):
                triplets = np.concatenate((triplets, np.array([(i, dt, dz, dz+1, dz+2)])))


    np.random.shuffle(triplets)
    triplets_train = triplets[:int((1-p_val)*len(triplets))]
    triplets_val = triplets[int((1-p_val)*len(triplets)):]

    if cfg.save_dataset :
        if os.path.exists('./dataset/'):
            shutil.rmtree('./dataset/')
        save_data(f'{root}/data/train/', f'{root}/dataset/train_{cfg.model_name}/', triplets_train)
        save_data(f'{root}/data/train/', f'{root}/dataset/val_{cfg.model_name}/', triplets_val)


def save_data(data_path, dataset_path, triplets):
    os.makedirs(dataset_path, exist_ok=True)

    files = os.listdir(data_path)
    for file_name in set(list(triplets[:,0])):
        selected_triplets = triplets[np.where(triplets[:,0] == file_name)]
        image = tiff.imread(f"{data_path}{files[file_name]}")

        format = 'y'
        len_shape = len(image.shape)
        if len_shape <= 2:
            raise ValueError('Error: images must be at least 3D.')
        elif len_shape == 3:
            image = image[None, :, None, :, :]
        elif len_shape == 4:
            if format == 'y':
                image = image[:, :, None, :, :]
            elif format == 'n':
                image = image[None, :, :, :, :]
            else:
                raise ValueError('Error: data file-format not recognized.')
        else:
            raise ValueError('Error: images must be at most 5D.')

        image = torch.from_numpy(image)  #.to(torch.uint16)
        image = image.to(torch.int32) # for torch.uint16 problem nuclei
        for tri in selected_triplets:
            torch.save(image[tri[1], tri[2:]], f"{dataset_path}{tri[0]}_{tri[1]}_{tri[2]}_{tri[3]}_{tri[4]}.pt")