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]