predictor.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  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 numpy as np
  12. from ...object_detection import DetPredictor
  13. from .keys import InstanceSegKeys as K
  14. from ..model_list import MODELS
  15. class InstanceSegPredictor(DetPredictor):
  16. """ Instance Seg Predictor """
  17. entities = MODELS
  18. def _run(self, batch_input):
  19. """ run """
  20. input_dict = {}
  21. input_dict["image"] = np.stack(
  22. [data[K.IMAGE] for data in batch_input], axis=0).astype(
  23. dtype=np.float32, copy=False)
  24. input_dict["scale_factor"] = np.stack(
  25. [data[K.SCALE_FACTOR][::-1] for data in batch_input],
  26. axis=0).astype(
  27. dtype=np.float32, copy=False)
  28. input_dict["im_shape"] = np.stack(
  29. [data[K.IM_SIZE][::-1] for data in batch_input], axis=0).astype(
  30. dtype=np.float32, copy=False)
  31. input_ = [input_dict[i] for i in self._predictor.get_input_names()]
  32. batch_np_boxes, batch_np_boxes_num, batch_np_masks = self._predictor.predict(
  33. input_)
  34. pred = batch_input
  35. box_idx_start = 0
  36. for idx in range(len(batch_input)):
  37. np_boxes_num = batch_np_boxes_num[idx]
  38. box_idx_end = box_idx_start + np_boxes_num
  39. np_boxes = batch_np_boxes[box_idx_start:box_idx_end]
  40. np_masks = batch_np_masks[box_idx_start:box_idx_end]
  41. box_idx_start = box_idx_end
  42. batch_input[idx][K.BOXES] = np_boxes
  43. batch_input[idx][K.MASKS] = np_masks
  44. return pred
  45. @classmethod
  46. def get_output_keys(cls):
  47. """ get output keys """
  48. return [K.BOXES, K.MASKS]