59 lines
1.4 KiB
Python
59 lines
1.4 KiB
Python
from typing import NewType
|
|
from dataclasses import dataclass
|
|
import numpy as np
|
|
import pandera as pa
|
|
from pandera.dtypes import Float32
|
|
|
|
XCoord = NewType('XCoord', float)
|
|
YCoord = NewType('YCoord', float)
|
|
Score = NewType('Confidence', float)
|
|
Category = NewType('Category', float)
|
|
|
|
|
|
|
|
@dataclass
|
|
class Point:
|
|
x: XCoord
|
|
y: YCoord
|
|
|
|
VideoFrame = NewType('Frame', float)
|
|
VideoTimestamp = float
|
|
|
|
@dataclass
|
|
class DetectedObject:
|
|
xmin: XCoord
|
|
ymin: YCoord
|
|
xmax: XCoord
|
|
ymax: YCoord
|
|
score: Score
|
|
category: Category
|
|
|
|
def as_int_box(self) -> tuple[tuple[int, int], tuple[int, int]]:
|
|
return (int(self.xmin), int(self.ymin)), (int(self.xmax), int(self.ymax))
|
|
|
|
def as_nparray(self):
|
|
return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score, self.category])
|
|
|
|
|
|
DetectedObject_ = pa.DataFrameSchema({
|
|
"xmin": pa.Column(Float32),
|
|
"ymin": pa.Column(Float32),
|
|
"xmax": pa.Column(Float32),
|
|
"ymin": pa.Column(Float32),
|
|
"score": pa.Column(Float32),
|
|
"category": pa.Column(Float32),
|
|
})
|
|
|
|
TrackedDetectedObject_ = pa.DataFrameSchema({
|
|
"xmin": pa.Column(Float32),
|
|
"ymin": pa.Column(Float32),
|
|
"xmax": pa.Column(Float32),
|
|
"ymin": pa.Column(Float32),
|
|
"object_id": pa.Column(Float32),
|
|
"score": pa.Column(Float32),
|
|
"category": pa.Column(Float32),
|
|
"index": pa.Column(Float32),
|
|
})
|
|
|
|
|
|
Color = tuple[int, int, int] |