| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556 |
- # !/usr/bin/env python3
- # -*- coding: UTF-8 -*-
- ################################################################################
- #
- # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
- #
- ################################################################################
- """
- Author: PaddlePaddle Authors
- """
- import numpy as np
- from ...object_detection import DetPredictor
- from .keys import InstanceSegKeys as K
- from ..model_list import MODELS
- class InstanceSegPredictor(DetPredictor):
- """ Instance Seg Predictor """
- entities = MODELS
- def _run(self, batch_input):
- """ run """
- input_dict = {}
- input_dict["image"] = np.stack(
- [data[K.IMAGE] for data in batch_input], axis=0).astype(
- dtype=np.float32, copy=False)
- input_dict["scale_factor"] = np.stack(
- [data[K.SCALE_FACTOR][::-1] for data in batch_input],
- axis=0).astype(
- dtype=np.float32, copy=False)
- input_dict["im_shape"] = np.stack(
- [data[K.IM_SIZE][::-1] for data in batch_input], axis=0).astype(
- dtype=np.float32, copy=False)
- input_ = [input_dict[i] for i in self._predictor.get_input_names()]
- batch_np_boxes, batch_np_boxes_num, batch_np_masks = self._predictor.predict(
- input_)
- pred = batch_input
- box_idx_start = 0
- for idx in range(len(batch_input)):
- np_boxes_num = batch_np_boxes_num[idx]
- box_idx_end = box_idx_start + np_boxes_num
- np_boxes = batch_np_boxes[box_idx_start:box_idx_end]
- np_masks = batch_np_masks[box_idx_start:box_idx_end]
- box_idx_start = box_idx_end
- batch_input[idx][K.BOXES] = np_boxes
- batch_input[idx][K.MASKS] = np_masks
- return pred
- @classmethod
- def get_output_keys(cls):
- """ get output keys """
- return [K.BOXES, K.MASKS]
|