| |
| |
| |
| import copy |
| from paddle.io import Dataset |
| from data_loader.modules import * |
|
|
|
|
| class BaseDataSet(Dataset): |
| def __init__( |
| self, |
| data_path: str, |
| img_mode, |
| pre_processes, |
| filter_keys, |
| ignore_tags, |
| transform=None, |
| target_transform=None, |
| ): |
| assert img_mode in ["RGB", "BRG", "GRAY"] |
| self.ignore_tags = ignore_tags |
| self.data_list = self.load_data(data_path) |
| item_keys = ["img_path", "img_name", "text_polys", "texts", "ignore_tags"] |
| for item in item_keys: |
| assert ( |
| item in self.data_list[0] |
| ), "data_list from load_data must contains {}".format(item_keys) |
| self.img_mode = img_mode |
| self.filter_keys = filter_keys |
| self.transform = transform |
| self.target_transform = target_transform |
| self._init_pre_processes(pre_processes) |
|
|
| def _init_pre_processes(self, pre_processes): |
| self.aug = [] |
| if pre_processes is not None: |
| for aug in pre_processes: |
| if "args" not in aug: |
| args = {} |
| else: |
| args = aug["args"] |
| if isinstance(args, dict): |
| cls = eval(aug["type"])(**args) |
| else: |
| cls = eval(aug["type"])(args) |
| self.aug.append(cls) |
|
|
| def load_data(self, data_path: str) -> list: |
| """ |
| 把数据加载为一个list： |
| :params data_path: 存储数据的文件夹或者文件 |
| return a dict ,包含了，'img_path','img_name','text_polys','texts','ignore_tags' |
| """ |
| raise NotImplementedError |
|
|
| def apply_pre_processes(self, data): |
| for aug in self.aug: |
| data = aug(data) |
| return data |
|
|
| def __getitem__(self, index): |
| try: |
| data = copy.deepcopy(self.data_list[index]) |
| im = cv2.imread(data["img_path"], 1 if self.img_mode != "GRAY" else 0) |
| if self.img_mode == "RGB": |
| im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) |
| data["img"] = im |
| data["shape"] = [im.shape[0], im.shape[1]] |
| data = self.apply_pre_processes(data) |
|
|
| if self.transform: |
| data["img"] = self.transform(data["img"]) |
| data["text_polys"] = data["text_polys"].tolist() |
| if len(self.filter_keys): |
| data_dict = {} |
| for k, v in data.items(): |
| if k not in self.filter_keys: |
| data_dict[k] = v |
| return data_dict |
| else: |
| return data |
| except: |
| return self.__getitem__(np.random.randint(self.__len__())) |
|
|
| def __len__(self): |
| return len(self.data_list) |
|
|
