cleaned up code.

dodgeballtracker/typing.py

@@ -2,7 +2,7 @@ from typing import NewType
 from dataclasses import dataclass
 import numpy as np
 import pandera as pa
-from pandera.dtypes import Float32
+from pandera.dtypes import Float32, Float64
 
 XCoord = NewType('XCoord', float)
 YCoord = NewType('YCoord', float)
@@ -10,15 +10,16 @@ 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
@@ -29,31 +30,24 @@ class DetectedObject:
     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))
+        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])
+        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),
-})
+def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema:
+    return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols})
 
 
-Color = tuple[int, int, int]
+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]