code_test/transformer/preprocess.py

import numpy as np
import cv2
import numbers
import logging

class Filter:
    # 根据需要添加其它函数
    def imnormalize(self, img, mean, std, to_rgb=True):
        """Inplace normalize an image with mean and std.
    
        Args:
            img (ndarray): Image to be normalized.
            mean (ndarray): The mean to be used for normalize.
            std (ndarray): The std to be used for normalize.
            to_rgb (bool): Whether to convert to rgb.
    
        Returns:
            ndarray: The normalized image.
        """
        # cv2 inplace normalization does not accept uint8
        #assert img.dtype != np.uint8
        mean = np.float64(mean.reshape(1, -1))
        stdinv = 1 / np.float64(std.reshape(1, -1))
        if to_rgb:
            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
        cv2.subtract(img, mean, img)  # inplace
        cv2.multiply(img, stdinv, img)  # inplace
        return img
        
    def impad(self, img,
          *,
          shape=None,
          padding=None,
          pad_val=0,
          padding_mode='constant'):
        assert (shape is not None) ^ (padding is not None)
        if shape is not None:
            padding = (0, 0, shape[1] - img.shape[1], shape[0] - img.shape[0])

        # check pad_val
        if isinstance(pad_val, tuple):
            assert len(pad_val) == img.shape[-1]
        elif not isinstance(pad_val, numbers.Number):
            raise TypeError('pad_val must be a int or a tuple. '
                            f'But received {type(pad_val)}')

        # check padding
        if isinstance(padding, tuple) and len(padding) in [2, 4]:
            if len(padding) == 2:
                padding = (padding[0], padding[1], padding[0], padding[1])
        elif isinstance(padding, numbers.Number):
            padding = (padding, padding, padding, padding)
        else:
            raise ValueError('Padding must be a int or a 2, or 4 element tuple.'
                             f'But received {padding}')

        # check padding mode
        assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']

        border_type = {
            'constant': cv2.BORDER_CONSTANT,
            'edge': cv2.BORDER_REPLICATE,
            'reflect': cv2.BORDER_REFLECT_101,
            'symmetric': cv2.BORDER_REFLECT
        }
        logging.error("padding...")
        logging.error(padding)
        self.pad_x = padding[0]
        self.pad_y = padding[1]
        img = cv2.copyMakeBorder(
            img,
            padding[1],
            padding[3],
            padding[0],
            padding[2],
            border_type[padding_mode],
            value=pad_val)

        return img

    # 可选提供init函数
    def init(self, params):
        self.means = params['normal.means']
        self.stds = params['normal.stds']
        self.resize_width = params['resize.width']
        self.resize_height = params['resize.height']

    # 必须提供filter函数
    def filter(self, inputs, meta_list):
        images_data = inputs['INPUT']
        images = []
        for image_data, meta in zip(images_data, meta_list):
            image_data = np.asarray(bytearray(image_data[0]), dtype="uint8")
            image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)

            h, w, _ = image.shape
            img_scale = (self.resize_width, self.resize_height)
            max_long_edge = max(img_scale)
            max_short_edge = min(img_scale)
            scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
            meta['scale_factor'] = scale_factor
            scale_w = int(w * float(scale_factor) + 0.5)
            scale_h = int(h * float(scale_factor) + 0.5)
            image = cv2.resize(image, (scale_w, scale_h))

            image = np.asarray(image).astype(np.float32)
            mean = np.array(self.means, dtype=np.float32)
            std = np.array(self.stds, dtype=np.float32)
            image = self.imnormalize(image, mean, std)

            divisor = 32
            pad_h = int(np.ceil(image.shape[0] / divisor)) * divisor
            pad_w = int(np.ceil(image.shape[1] / divisor)) * divisor
            image = self.impad(image, shape=(pad_h, pad_w), pad_val=0)
            image = np.asarray(image).astype(np.float32)
            image = np.transpose(image, [2, 0, 1])

            images.append(image)

        return {'input': images}

    # 可选提供finalize函数
    def finalize(self):
        pass