config.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. from ....utils.misc import abspath
  13. from ..text_rec.config import TextRecConfig
  14. class TextDetConfig(TextRecConfig):
  15. """ Text Detection Config """
  16. def update_batch_size(self, batch_size: int):
  17. """update batch size setting
  18. Args:
  19. batch_size (int): the batch size number of training loader to set.
  20. """
  21. _cfg = {'Train.loader.batch_size_per_card': batch_size, }
  22. self.update(_cfg)
  23. def update_dataset(self, dataset_path: str, dataset_type=None):
  24. """update dataset settings
  25. Args:
  26. dataset_path (str): the root path of dataset.
  27. dataset_type (str, optional): dataset type. Defaults to None.
  28. Raises:
  29. ValueError: the dataset_type error.
  30. """
  31. dataset_path = abspath(dataset_path)
  32. if dataset_type is None:
  33. dataset_type = 'TextDetDataset'
  34. if dataset_type == 'TextDetDataset':
  35. _cfg = {
  36. 'Train.dataset.name': dataset_type,
  37. 'Train.dataset.data_dir': dataset_path,
  38. 'Train.dataset.label_file_list':
  39. [os.path.join(dataset_path, 'train.txt')],
  40. 'Eval.dataset.name': dataset_type,
  41. 'Eval.dataset.data_dir': dataset_path,
  42. 'Eval.dataset.label_file_list':
  43. [os.path.join(dataset_path, 'val.txt')]
  44. }
  45. self.update(_cfg)
  46. else:
  47. raise ValueError(f"{repr(dataset_type)} is not supported.")