Deep learning (target tracking and target detection) -- bounding box bbox coordinate conversion (arbitrary format [list, numpy, tensor], arbitrary dimension [vector, one-dimensional matrix, two-dimensional matrix])

Note from the author: There may be errors, and it can run on my computer;
during the process of writing the program, I found that the bounding box conversion programs written by different people are different,
some can only convert numpy matrices, and
some can only convert tensor matrices
. I tried to write a bbox function that can convert any format of any dimension . The level is not enough, it took a long time to write, and my head was a little dizzy, so I sent it out, hoping that everyone can find the mistakes in it, and it is also convenient for everyone to use; if friends find problems with the program, I hope you can point it out in time, I Amendments will be made immediately to make progress together.

The purpose of this program is to convert the input data list, numpy, and tensor in the following three formats, and the dimensions can range from 0 to 2 dimensions, that is, the shape is: (4,) (3, 4) torch.Size([4] ) bounding box data of torch.Size([3, 4])

import numpy as np
import torch


#  ===============================================================================#
#  坐标转换系列函数
#  输入:可能是 列表、np矩阵、tensor矩阵 以下六个函数可以保证输入输出的维度一致
#  输入的维度可能是一个向量shape=(4,)(.T转置之后的到的是原变量)
#  ===============================================================================#
def ltwh2center(bbox):
    """

    :param bbox:[left, top, w, h]
    :return:[cx, cy, w, h]
    """
    if isinstance(bbox, list):
        bbox = np.array(bbox)

    if bbox.shape[-1] != 4:
        raise ValueError('bbox.shape[-1] should equal 4')
    else:
        if isinstance(bbox, np.ndarray):
            left, top, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # cx=left+w/2; cy=top+h/2;w;h
            _bbox = np.stack([left + w / 2, top + h / 2, w, h], axis=-1)
            return _bbox

        if isinstance(bbox, torch.Tensor):
            left, top, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # cx=left+w/2; cy=top+h/2;w;h
            _bbox = torch.stack((left + w / 2, top + h / 2, w, h), dim=-1)
            return _bbox


def ltwh2corner(bbox):
    """

    :param bbox:[left, top, w, h]
    :return:[left, top, right, bottom]
    """
    if isinstance(bbox, list):
        bbox = np.array(bbox)

    if bbox.shape[-1] != 4:
        raise ValueError('bbox.shape[-1] should equal 4')
    else:
        if isinstance(bbox, np.ndarray):
            left, top, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # left; top; right=left+w; bottom=top+h
            _bbox = np.stack([left, top, left + w, top + h], axis=-1)
            return _bbox

        if isinstance(bbox, torch.Tensor):
            left, top, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            _bbox = torch.stack((left, top, left + w, top + h), dim=-1)
            return _bbox


def corner2ltwh(bbox):
    """

    :param bbox:[left, top, right, bottom]
    :return:[left, top, w, h]
    """
    if isinstance(bbox, list):
        bbox = np.array(bbox)

    if bbox.shape[-1] != 4:
        raise ValueError('bbox.shape[-1] should equal 4')
    else:
        if isinstance(bbox, np.ndarray):
            left, top, right, bottom = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # left; top; w=right-left; h=bottom-top
            _bbox = np.stack([left, top, right - left, bottom - top], axis=-1)
            return _bbox

        if isinstance(bbox, torch.Tensor):
            left, top, right, bottom = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            _bbox = torch.stack((left, top, right - left, bottom - top), dim=-1)
            return _bbox


def corner2center(bbox):
    """

    :param bbox:[left, top, right, bottom]
    :return:[cx,cy, w, h]
    """
    if isinstance(bbox, list):
        bbox = np.array(bbox)

    if bbox.shape[-1] != 4:
        raise ValueError('bbox.shape[-1] should equal 4')
    else:
        if isinstance(bbox, np.ndarray):
            left, top, right, bottom = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # cx=(left+right)/2; cy=(top+bottom)/2; w=right-left; h=bottom-top
            _bbox = np.stack([(left + right) / 2, (top + bottom) / 2, right - left, bottom - top], axis=-1)
            return _bbox

        if isinstance(bbox, torch.Tensor):
            left, top, right, bottom = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            _bbox = torch.stack(((left + right) / 2, (top + bottom) / 2, right - left, bottom - top), dim=-1)
            return _bbox


def center2corner(bbox):
    """

    :param bbox: [cx,cy,w,h]
    :return: [left, top, right, bottom]
    """
    if isinstance(bbox, list):
        bbox = np.array(bbox)

    if bbox.shape[-1] != 4:
        raise ValueError('bbox.shape[-1] should equal 4')
    else:
        if isinstance(bbox, np.ndarray):
            cx, cy, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # left=cx-w/2; top=cy-h/2; right=cx+w/2; bottom=cy+h/2
            _bbox = np.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], axis=-1)
            return _bbox

        if isinstance(bbox, torch.Tensor):
            cx, cy, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            _bbox = torch.stack((cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2), dim=-1)
            return _bbox


def center2ltwh(bbox):
    """

    :param bbox: [cx, cy, w, h]
    :return: [left, top, w, h]
    """
    if isinstance(bbox, list):
        bbox = np.array(bbox)

    if bbox.shape[-1] != 4:
        raise ValueError('bbox.shape[-1] should equal 4')
    else:
        if isinstance(bbox, np.ndarray):
            cx, cy, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            # left=cx-w/2; top=cy-h/2; w; h
            _bbox = np.stack([cx - w / 2, cy - h / 2, w, h], axis=-1)  # cx,cy,w,h
            return _bbox

        if isinstance(bbox, torch.Tensor):
            cx, cy, w, h = bbox[..., 0], bbox[..., 1], bbox[..., 2], bbox[..., 3]
            _bbox = torch.stack((cx - w / 2, cy - h / 2, w, h), dim=-1)  # 将数据坐标拼接起来
            return _bbox


if __name__ == '__main__':
    print('Start...')
    box1 = [50, 50, 100, 200]  # list
    box2 = np.array([50, 50, 120, 220])  # 一个坐标
    box3 = np.array([[50, 50, 100, 200], [50, 50, 120, 220], [50, 50, 120, 220]])  # 多个坐标
    box4 = torch.FloatTensor([50, 50, 100, 200])  # 一个tensor坐标数据
    box5 = torch.FloatTensor([[50, 50, 100, 200], [50, 50, 120, 220], [50, 50, 120, 220]])  # 多个tensor坐标数据

    for box in [box1, box2, box3, box4, box5]:
        box_ = ltwh2center(box)
        print('\n', 'input (%s):\n' % type(box), box, '\n', 'output(%s):\n' % type(box_), box_)

Guess you like

Origin blog.csdn.net/weixin_50727642/article/details/122737638