32 lines
689 B
Python
32 lines
689 B
Python
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
|
|
@dataclass
|
|
class DetectionResult:
|
|
bbox: List[int] # [x1, y1, x2, y2]
|
|
class_id: int
|
|
class_name: str
|
|
confidence: float
|
|
track_id: int
|
|
|
|
|
|
@dataclass
|
|
class DetectionResultList:
|
|
boxes: [] # [x1, y1, x2, y2]
|
|
clss: []
|
|
clss_name: []
|
|
confs: []
|
|
track_ids: []
|
|
|
|
def __iter__(self):
|
|
"""迭代返回每个检测结果的 (bbox, cls_id, cls_name, conf)"""
|
|
for i in range(len(self.boxes)):
|
|
yield (
|
|
self.boxes[i],
|
|
self.clss[i],
|
|
self.clss_name[i],
|
|
self.confs[i],
|
|
self.track_ids[i]
|
|
)
|