multilabel_classification.py 4.6 KB

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