54 lines
1.4 KiB
Python
54 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, Float64
|
|
|
|
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])
|
|
|
|
|
|
def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema:
|
|
return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols})
|
|
|
|
|
|
DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score",
|
|
"category"])
|
|
|
|
TrackedDetectedObject_ = tensor_of(Float64,
|
|
list(DetectedObject_.columns) + ["object_id", "index"])
|
|
# ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"])
|
|
|
|
|
|
Color = tuple[int, int, int]
|