回传数据解析,兼容v5和v10
This commit is contained in:
95
tracking/utils/proBoxes.py
Normal file
95
tracking/utils/proBoxes.py
Normal file
@ -0,0 +1,95 @@
|
||||
# -*- 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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
Reference in New Issue
Block a user