目标检测中数据预处理

0 前言

在计算机视觉任务中，大部分神经网络要求输入图像数据的高、宽一致。尤其是在目标检测任务中，原始图像在resize后，还需要对原始坐标信息进行重映射。
本文要做的是 填充+resize+重映射

1 代码

import numpy as np
from PIL import Image


def preprocess_image(image, size=(416,416)):
    '''
    对图片进行 灰度填充 + resize
    :param image: 传入一张尺寸不定的图片
    :param size: SSD输入的尺寸
    :return: img, 宽度偏移量，高度偏移量，max(w,h)
    '''

    iw, ih = image.size
    max_length = max(iw, ih)
    dw = (max_length - iw) // 2
    dh = (max_length - ih) // 2
    if iw != ih:
        new_image = Image.new(mode='RGB', size=(max_length, max_length), color=(128, 128, 128))
        new_image.paste(im=image, box=(dw, dh))
        new_image = new_image.resize(size, Image.BICUBIC)
    else:
        new_image = image.resize(size, Image.BICUBIC)
    return np.asarray(new_image), dw, dh, max_length


def preprocess_box(bboxes, dw, dh, max_length, size):
    # 将坐标信息映射在新图中
    pre_bbox = np.zeros(bboxes.shape, dtype=np.float32)
    pre_bbox[:, [0,2]] = bboxes[:, [0,2]] + dw
    pre_bbox[:, [1,3]] = bboxes[:, [1,3]] + dh
    pre_bbox[:,4] = bboxes[:, 4]
    pre_bbox[:,:4] /= max_length
    pre_bbox[:,:4] *= size[0]
    return pre_bbox


def preprocess_annotation(image, bboxes, size):
    '''
    功能：数据预处理，将图片进行填充，并resize到高、宽相等的尺寸。同时对坐标信息重新映射
    :param image: 输入图片，必须是由PIL.Image.open()读入的
    :param bboxes: 输入坐标信息，np.array格式[[xmin,ymin,xmax,ymax,lab], [xmin,ymin,xmax,ymax,lab], ...]
    :param size: 网络输入尺寸，即图片resize的尺寸，元组(int, int)
    :return: 预处理的图片和坐标信息
    '''
    pre_image, dw, dh, max_length = preprocess_image(image, size)
    pre_bbox = preprocess_box(bboxes, dw, dh, max_length, size)
    return pre_image, pre_bbox

使用方法： 只需要调用 preprocess_annotation(image, bboxes, size) 函数，传入待处理的图像、坐标以及网络输入尺寸。