diff --git a/ppocr/data/simple_dataset.py b/ppocr/data/simple_dataset.py index 969260b..e470d46 100644 --- a/ppocr/data/simple_dataset.py +++ b/ppocr/data/simple_dataset.py @@ -116,42 +116,57 @@ class SimpleDataSet(Dataset): ext_data.append(data) return ext_data + def create_dummy_sample(self): + """Создает пустой образец для предотвращения сбоев""" + dummy_img = np.zeros((32, 320, 3), dtype=np.uint8) + dummy_img = cv2.imencode('.jpg', dummy_img)[1].tobytes() + + return { + "image": dummy_img, + "label": "", + "img_path": "dummy.jpg", + "ext_data": self.get_ext_data(), + "filename": "dummy.jpg" + } + def __getitem__(self, idx): - file_idx = self.data_idx_order_list[idx] - data_line = self.data_lines[file_idx] - try: - data_line = data_line.decode("utf-8") - substr = data_line.strip("\n").split(self.delimiter) - file_name = substr[0] - file_name = self._try_parse_filename_list(file_name) - 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)) - with open(data["img_path"], "rb") as f: - img = f.read() - data["image"] = img - data["ext_data"] = self.get_ext_data() - data["filename"] = data["img_path"] - outs = transform(data, self.ops) - except Exception as exc: - print("HERE", str(exc)) - self.logger.error( - "When parsing line {}, error happened with msg: {}".format( - data_line, traceback.format_exc() + max_attempts = 10 + for attempt in range(max_attempts): + try: + file_idx = self.data_idx_order_list[idx] + 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] + file_name = self._try_parse_filename_list(file_name) + 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)) + with open(data["img_path"], "rb") as f: + img = f.read() + data["image"] = img + data["ext_data"] = self.get_ext_data() + data["filename"] = data["img_path"] + outs = transform(data, self.ops) + return outs + + except Exception as exc: + if attempt < max_attempts - 1: + idx = ( + np.random.randint(self.__len__()) + if self.mode == "train" + else (idx + 1) % self.__len__() + ) + + self.logger.error( + "Failed after {} attempts: {}".format( + max_attempts, str(exc) + ) ) - ) - outs = None - if outs is None: - # 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) - return outs + + return self.create_dummy_sample() def __len__(self): return len(self.data_idx_order_list)