multilabel_classification.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  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. from typing import Any, Dict, List, Optional, Union
  16. import ultra_infer as ui
  17. import numpy as np
  18. from pathlib import Path
  19. import tempfile
  20. import yaml
  21. from paddlex.inference.common.batch_sampler import ImageBatchSampler
  22. from paddlex.inference.models.image_multilabel_classification.result import (
  23. MLClassResult,
  24. )
  25. from paddlex.modules.multilabel_classification.model_list import MODELS
  26. from paddlex.utils import logging
  27. from paddlex_hpi.models.base import CVPredictor, HPIParams
  28. class MLClasPredictor(CVPredictor):
  29. entities = MODELS
  30. def __init__(
  31. self,
  32. model_dir: Union[str, os.PathLike],
  33. config: Optional[Dict[str, Any]] = None,
  34. device: Optional[str] = None,
  35. batch_size: int = 1,
  36. hpi_params: Optional[HPIParams] = None,
  37. threshold: Union[float, dict, list, None] = None,
  38. ) -> None:
  39. self._threshold = threshold
  40. super().__init__(
  41. model_dir=model_dir,
  42. config=config,
  43. device=device,
  44. batch_size=batch_size,
  45. hpi_params=hpi_params,
  46. )
  47. self._label_list = self._get_label_list()
  48. def _build_ui_model(
  49. self, option: ui.RuntimeOption
  50. ) -> ui.vision.classification.PyOnlyMultilabelClassificationModel:
  51. if self._threshold:
  52. if not isinstance(self._threshold, (float, None)):
  53. logging.warning("`threshold` must be float or None in PaddleX HPI")
  54. with open(self.config_path, "r") as file:
  55. config = yaml.safe_load(file)
  56. config["PostProcess"]["MultiLabelThreshOutput"][
  57. "threshold"
  58. ] = self._threshold
  59. temp_dir = os.path.dirname(self.config_path)
  60. with tempfile.NamedTemporaryFile(
  61. delete=False, dir=temp_dir, suffix=".yml", mode="w", encoding="utf-8"
  62. ) as temp_file:
  63. temp_file_path = temp_file.name
  64. yaml.safe_dump(config, temp_file, default_flow_style=False)
  65. model = ui.vision.classification.PyOnlyMultilabelClassificationModel(
  66. str(self.model_path),
  67. str(self.params_path),
  68. str(Path(temp_file_path)),
  69. runtime_option=option,
  70. )
  71. else:
  72. model = ui.vision.classification.PyOnlyMultilabelClassificationModel(
  73. str(self.model_path),
  74. str(self.params_path),
  75. str(self.config_path),
  76. runtime_option=option,
  77. )
  78. return model
  79. def _build_batch_sampler(self) -> ImageBatchSampler:
  80. return ImageBatchSampler()
  81. def _get_result_class(self) -> type:
  82. return MLClassResult
  83. def process(
  84. self,
  85. batch_data: List[Any],
  86. threshold: Union[float, dict, list, None] = None,
  87. ) -> Dict[str, List[Any]]:
  88. if threshold:
  89. logging.warning(
  90. "`threshold` is not supported for multilabel classification in PaddleX HPI"
  91. )
  92. batch_raw_imgs = self._data_reader(imgs=batch_data.instances)
  93. imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
  94. ui_results = self._ui_model.batch_predict(imgs)
  95. class_ids_list = []
  96. scores_list = []
  97. label_names_list = []
  98. for ui_result in ui_results:
  99. class_ids_list.append(ui_result.label_ids)
  100. scores_list.append(np.around(ui_result.scores, decimals=5).tolist())
  101. if self._label_list is not None:
  102. label_names_list.append(
  103. [self._label_list[i] for i in ui_result.label_ids]
  104. )
  105. return {
  106. "input_path": batch_data.input_paths,
  107. "page_index": batch_data.page_indexes,
  108. "input_img": batch_raw_imgs,
  109. "class_ids": class_ids_list,
  110. "scores": scores_list,
  111. "label_names": label_names_list,
  112. }
  113. def _get_label_list(self) -> Optional[List[str]]:
  114. pp_config = self.config["PostProcess"]
  115. if "MultiLabelThreshOutput" not in pp_config:
  116. raise RuntimeError("`MultiLabelThreshOutput` config not found")
  117. label_list = pp_config["MultiLabelThreshOutput"].get("label_list", None)
  118. return label_list