initial project version!

tracking/utils/proBoxes.py

@@ -0,0 +1,94 @@
+# -*- 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, :]
+        n = boxes.shape[-1]
+        assert n in (6, 7, 8), f'expected `n` in [6, 7], but got {n}'  # xyxyb, track_id, conf, cls
+        
+        self.data = boxes
+        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
+
+
+
+
+
+
+
+
+
+
+
+