96 lines
2.4 KiB
Python
96 lines
2.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Created on Sun Dec 31 17:14:37 2023
|
||
|
||
@author: ym
|
||
"""
|
||
import numpy as np
|
||
|
||
|
||
class Boxes:
|
||
def __init__(self, boxes, orig_shape=None) -> None:
|
||
"""Initialize the Boxes class."""
|
||
if boxes.ndim == 1:
|
||
boxes = boxes[None, :]
|
||
m, n = boxes.shape
|
||
assert n in (6, 7), f'expected `n` in [6, 7], but got {n}' # xyxy, track_id, conf, cls
|
||
|
||
'''对每一个box进行编号，利用该编号可以索引对应 feature'''
|
||
self.data = np.concatenate([boxes[:, :4], np.arange(m).reshape(-1, 1), boxes[:, 4:]], axis=-1)
|
||
|
||
self.orig_shape = orig_shape
|
||
|
||
def cpu(self):
|
||
"""Return a copy of the tensor on CPU memory."""
|
||
return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.cpu(), self.orig_shape)
|
||
def numpy(self):
|
||
"""Return a copy of the tensor as a numpy array."""
|
||
return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.numpy(), self.orig_shape)
|
||
|
||
@property
|
||
def xyxy(self):
|
||
"""Return the boxes in xyxy format."""
|
||
return self.data[:, :4]
|
||
|
||
@property
|
||
def xyxyb(self):
|
||
"""Return the boxes in xyxyb format."""
|
||
return self.data[:, :5]
|
||
|
||
@property
|
||
def conf(self):
|
||
"""Return the confidence values of the boxes."""
|
||
return self.data[:, -2]
|
||
|
||
@property
|
||
def cls(self):
|
||
"""Return the class values of the boxes."""
|
||
return self.data[:, -1]
|
||
|
||
|
||
|
||
# def boxes_add_fid(tboxes):
|
||
# '''
|
||
# 将 bboxes 对应的帧索引添加到 boxes 最后一列
|
||
# Return：
|
||
# bboxes: [x1, y1, x2, y2, track_id, score, cls, frame_index]
|
||
# '''
|
||
|
||
# bboxes = np.empty((0, 8), dtype = np.float32)
|
||
# for tbox, f in tboxes:
|
||
# data = tbox.numpy()
|
||
|
||
# frame = f * np.ones([data.shape[0], 1])
|
||
# bbox = np.concatenate([data, frame], axis=1)
|
||
# bboxes = np.concatenate([bboxes, bbox], axis=0)
|
||
|
||
# return bboxes
|
||
|
||
|
||
def boxes_add_fid(tboxes):
|
||
'''
|
||
将 bboxes 对应的帧索引添加到 boxes 最后一列
|
||
Return：
|
||
bboxes: [x1, y1, x2, y2, track_id, score, cls, frame_index]
|
||
'''
|
||
|
||
bboxes = np.empty((0, 8), dtype = np.float32)
|
||
for data, f in tboxes:
|
||
frame = f * np.ones([data.shape[0], 1])
|
||
bbox = np.concatenate([data, frame], axis=1)
|
||
bboxes = np.concatenate([bboxes, bbox], axis=0)
|
||
|
||
return bboxes
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|