image-segmentation/data.py at master · zhangjingcode/image-segmentation · 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
import h5py
import numpy as np
import os


def GetData(data_folder):
    file_list = os.listdir(data_folder)
    image_list = []
    label_list = []

    for file in file_list:
        file_path = os.path.join(data_folder, file)

        # data read
        h5_file = h5py.File(file_path, 'r')
        image = np.asarray(h5_file['input_0'], dtype=np.float32)
        label = np.asarray(h5_file['output_0'], dtype=np.uint8)
        h5_file.close()

        image_list.append(image)
        label_list.append(label)

    return np.asarray(image_list), np.asarray(label_list)


def GeneratorData(data_folder, batch_size):
    file_list = os.listdir(data_folder)
    image_list = []
    label_list = []

    while True:
        for file in file_list:
            file_path = os.path.join(data_folder, file)

            # data read
            h5_file = h5py.File(file_path, 'r')
            image = np.asarray(h5_file['input_0'], dtype=np.float32)
            label = np.asarray(h5_file['output_0'], dtype=np.uint8)
            h5_file.close()

            image_list.append(image)
            label_list.append(label)

            if len(image_list) >= batch_size:
                yield np.asarray(image_list), np.asarray(label_list)
                image_list = []
                label_list = []