simple_dataset.py 5.18 KB
Newer Older
dyning's avatar
dyning committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import os
import random
17
import traceback
dyning's avatar
dyning committed
18
19
from paddle.io import Dataset
from .imaug import transform, create_operators
dyning's avatar
dyning committed
20

dyning's avatar
dyning committed
21
22

class SimpleDataSet(Dataset):
23
    def __init__(self, config, mode, logger, seed=None):
dyning's avatar
dyning committed
24
        super(SimpleDataSet, self).__init__()
25
        self.logger = logger
26
        self.mode = mode.lower()
dyning's avatar
dyning committed
27

dyning's avatar
dyning committed
28
29
30
        global_config = config['Global']
        dataset_config = config[mode]['dataset']
        loader_config = config[mode]['loader']
dyning's avatar
dyning committed
31

dyning's avatar
dyning committed
32
33
34
        self.delimiter = dataset_config.get('delimiter', '\t')
        label_file_list = dataset_config.pop('label_file_list')
        data_source_num = len(label_file_list)
35
36
        ratio_list = dataset_config.get("ratio_list", [1.0])
        if isinstance(ratio_list, (float, int)):
LDOUBLEV's avatar
LDOUBLEV committed
37
            ratio_list = [float(ratio_list)] * int(data_source_num)
dyning's avatar
dyning committed
38
39
40
41

        assert len(
            ratio_list
        ) == data_source_num, "The length of ratio_list should be the same as the file_list."
dyning's avatar
dyning committed
42
43
        self.data_dir = dataset_config['data_dir']
        self.do_shuffle = loader_config['shuffle']
dyning's avatar
dyning committed
44

45
        self.seed = seed
dyning's avatar
dyning committed
46
        logger.info("Initialize indexs of datasets:%s" % label_file_list)
LDOUBLEV's avatar
LDOUBLEV committed
47
        self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
48
        self.data_idx_order_list = list(range(len(self.data_lines)))
49
        if self.mode == "train" and self.do_shuffle:
50
            self.shuffle_data_random()
dyning's avatar
dyning committed
51
52
        self.ops = create_operators(dataset_config['transforms'], global_config)

LDOUBLEV's avatar
LDOUBLEV committed
53
    def get_image_info_list(self, file_list, ratio_list):
dyning's avatar
dyning committed
54
55
        if isinstance(file_list, str):
            file_list = [file_list]
56
57
        data_lines = []
        for idx, file in enumerate(file_list):
dyning's avatar
dyning committed
58
59
            with open(file, "rb") as f:
                lines = f.readlines()
60
61
62
63
                if self.mode == "train" or ratio_list[idx] < 1.0:
                    random.seed(self.seed)
                    lines = random.sample(lines,
                                          round(len(lines) * ratio_list[idx]))
64
65
                data_lines.extend(lines)
        return data_lines
dyning's avatar
dyning committed
66
67

    def shuffle_data_random(self):
68
69
        random.seed(self.seed)
        random.shuffle(self.data_lines)
dyning's avatar
dyning committed
70
        return
dyning's avatar
dyning committed
71

WenmuZhou's avatar
WenmuZhou committed
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
    def get_ext_data(self):
        ext_data_num = 0
        for op in self.ops:
            if hasattr(op, 'ext_data_num'):
                ext_data_num = getattr(op, 'ext_data_num')
                break
        load_data_ops = self.ops[:2]
        ext_data = []

        while len(ext_data) < ext_data_num:
            file_idx = self.data_idx_order_list[np.random.randint(self.__len__(
            ))]
            data_line = self.data_lines[file_idx]
            data_line = data_line.decode('utf-8')
            substr = data_line.strip("\n").split(self.delimiter)
            file_name = substr[0]
            label = substr[1]
            img_path = os.path.join(self.data_dir, file_name)
            data = {'img_path': img_path, 'label': label}
            if not os.path.exists(img_path):
                continue
            with open(data['img_path'], 'rb') as f:
                img = f.read()
                data['image'] = img
            data = transform(data, load_data_ops)
97
98

            if data is None or data['polys'].shape[1]!=4:
WenmuZhou's avatar
WenmuZhou committed
99
100
101
102
                continue
            ext_data.append(data)
        return ext_data

dyning's avatar
dyning committed
103
    def __getitem__(self, idx):
104
        file_idx = self.data_idx_order_list[idx]
LDOUBLEV's avatar
LDOUBLEV committed
105
        data_line = self.data_lines[file_idx]
106
        try:
LDOUBLEV's avatar
LDOUBLEV committed
107
108
109
110
111
112
113
114
            data_line = data_line.decode('utf-8')
            substr = data_line.strip("\n").split(self.delimiter)
            file_name = substr[0]
            label = substr[1]
            img_path = os.path.join(self.data_dir, file_name)
            data = {'img_path': img_path, 'label': label}
            if not os.path.exists(img_path):
                raise Exception("{} does not exist!".format(img_path))
115
116
117
            with open(data['img_path'], 'rb') as f:
                img = f.read()
                data['image'] = img
WenmuZhou's avatar
WenmuZhou committed
118
            data['ext_data'] = self.get_ext_data()
119
            outs = transform(data, self.ops)
120
        except:
121
            self.logger.error(
LDOUBLEV's avatar
LDOUBLEV committed
122
                "When parsing line {}, error happened with msg: {}".format(
123
                    data_line, traceback.format_exc()))
124
            outs = None
dyning's avatar
dyning committed
125
        if outs is None:
126
127
128
129
            # during evaluation, we should fix the idx to get same results for many times of evaluation.
            rnd_idx = np.random.randint(self.__len__(
            )) if self.mode == "train" else (idx + 1) % self.__len__()
            return self.__getitem__(rnd_idx)
dyning's avatar
dyning committed
130
131
132
133
        return outs

    def __len__(self):
        return len(self.data_idx_order_list)