check_dataset.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import json
  16. import os.path as osp
  17. from collections import defaultdict
  18. from .....utils.errors import DatasetFileNotFoundError, CheckFailedError
  19. def check(dataset_dir, output, dataset_type="PubTabTableRecDataset", sample_num=10):
  20. """
  21. Check whether the dataset is valid.
  22. """
  23. if dataset_type == "PubTabTableRecDataset":
  24. # Custom dataset
  25. if not osp.exists(dataset_dir) or not osp.isdir(dataset_dir):
  26. raise DatasetFileNotFoundError(file_path=dataset_dir)
  27. tags = ["train", "val"]
  28. max_recorded_sample_cnts = 50
  29. sample_cnts = dict()
  30. sample_paths = defaultdict(list)
  31. for tag in tags:
  32. file_list = osp.join(dataset_dir, f"{tag}.txt")
  33. if not osp.exists(file_list):
  34. if tag in ("train", "val"):
  35. # train and val file lists must exist
  36. raise DatasetFileNotFoundError(
  37. file_path=file_list,
  38. solution=f"Ensure that both `train.txt` and `val.txt` exist in {dataset_dir}",
  39. )
  40. else:
  41. # tag == 'test'
  42. continue
  43. else:
  44. with open(file_list, "r", encoding="utf-8") as f:
  45. all_lines = f.readlines()
  46. sample_cnts[tag] = len(all_lines)
  47. for line in all_lines:
  48. info = json.loads(line.strip("\n"))
  49. file_name = info["filename"]
  50. cells = info["html"]["cells"].copy()
  51. structure = info["html"]["structure"]["tokens"].copy()
  52. img_path = osp.join(dataset_dir, file_name)
  53. if len(sample_paths[tag]) < max_recorded_sample_cnts:
  54. sample_paths[tag].append(os.path.relpath(img_path, output))
  55. if not os.path.exists(img_path):
  56. raise DatasetFileNotFoundError(file_path=img_path)
  57. boxes_num = len(cells)
  58. tokens_num = sum(
  59. [
  60. structure.count(x)
  61. for x in ["<td>", "<td", "<eb></eb>", "<td></td>"]
  62. ]
  63. )
  64. if boxes_num != tokens_num:
  65. raise CheckFailedError(
  66. f"The number of cells needs to be consistent with the number of tokens "
  67. "but the number of cells is {boxes_num}, and the number of tokens is {tokens_num}."
  68. )
  69. meta = {}
  70. meta["train_samples"] = sample_cnts["train"]
  71. meta["train_sample_paths"] = sample_paths["train"][:sample_num]
  72. meta["val_samples"] = sample_cnts["val"]
  73. meta["val_sample_paths"] = sample_paths["val"][:sample_num]
  74. return meta