| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778 |
- # !/usr/bin/env python3
- # -*- coding: UTF-8 -*-
- ################################################################################
- #
- # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
- #
- ################################################################################
- """
- Author: PaddlePaddle Authors
- """
- import os
- import numpy as np
- from pathlib import Path
- from ...base import BasePredictor
- from ...base.predictor.transforms import image_common
- from .keys import ClsKeys as K
- from .utils import InnerConfig
- from ....utils import logging
- from . import transforms as T
- from ..model_list import MODELS
- class ClsPredictor(BasePredictor):
- """ Clssification Predictor """
- entities = MODELS
- def load_other_src(self):
- """ load the inner config file """
- infer_cfg_file_path = os.path.join(self.model_dir, 'inference.yml')
- if not os.path.exists(infer_cfg_file_path):
- raise FileNotFoundError(
- f"Cannot find config file: {infer_cfg_file_path}")
- return InnerConfig(infer_cfg_file_path)
- @classmethod
- def get_input_keys(cls):
- """ get input keys """
- return [[K.IMAGE], [K.IM_PATH]]
- @classmethod
- def get_output_keys(cls):
- """ get output keys """
- return [K.CLS_PRED]
- def _run(self, batch_input):
- """ run """
- input_dict = {}
- input_dict[K.IMAGE] = np.stack(
- [data[K.IMAGE] for data in batch_input], axis=0).astype(
- dtype=np.float32, copy=False)
- input_ = [input_dict[K.IMAGE]]
- outputs = self._predictor.predict(input_)
- cls_outs = outputs[0]
- # In-place update
- pred = batch_input
- for dict_, cls_out in zip(pred, cls_outs):
- dict_[K.CLS_PRED] = cls_out
- return pred
- def _get_pre_transforms_from_config(self):
- """ get preprocess transforms """
- logging.info(
- f"Transformation operators for data preprocessing will be inferred from config file."
- )
- pre_transforms = self.other_src.pre_transforms
- pre_transforms.insert(0, image_common.ReadImage(format='RGB'))
- return pre_transforms
- def _get_post_transforms_from_config(self):
- """ get postprocess transforms """
- post_transforms = self.other_src.post_transforms
- post_transforms.extend([
- T.PrintResult(), T.SaveClsResults(self.output,
- self.other_src.labels)
- ])
- return post_transforms
|