21 lines
566 B
Python
21 lines
566 B
Python
from typing import NewType
|
|
import pandera as pa
|
|
from pandera.dtypes import Float32, Float64
|
|
|
|
VideoFrame = NewType('Frame', float)
|
|
VideoTimestamp = float
|
|
|
|
|
|
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"])
|
|
|
|
|
|
|
|
Color = tuple[int, int, int]
|